Полёт с удержанием ArUco-метки

from pioneer_sdk2 import Pioneer, Camera, ImageViewer # импортируем классы Pioneer, Camera, ImageViewer из библиотеки pioneer_sdk2
import cv2                                            # библиотека cv2 содержит функции для работы с изображениями
import numpy as np                                    # библиотека numpy нужна для работы с массивами
import threading                                      # библиотека threading нужна для запуска обработки видео в отдельном потоке
import time                                           # библиотека time содержит функции для работы со временем
import sys                                            # библиотека sys нужна для чтения команд из терминала
import select                                         # библиотека select нужна для неблокирующего чтения команд
import termios                                        # библиотека termios нужна для настройки терминала
import tty                                            # библиотека tty нужна для чтения клавиш без Enter
from pathlib import Path                              # класс Path нужен для работы с путем к файлу калибровки

try:                                                  # пробуем импортировать управление сервокамерой
    from pioneer_sdk2 import ServoCamera, ServoPriority # импортируем классы для управления углом камеры
except ImportError:                                   # если текущая конфигурация SDK не поддерживает сервокамеру
    ServoCamera = None                                # сохраняем None вместо класса ServoCamera
    ServoPriority = None                              # сохраняем None вместо класса ServoPriority


DATA_PATH = Path(__file__).with_name("data.yml")       # путь к файлу калибровки рядом со скриптом
DEBUG_PRINT_INTERVAL = 0.5                             # пауза между диагностическими сообщениями в терминал
CAMERA_HORIZONTAL_ANGLE = 25                           # угол камеры, при котором она смотрит горизонтально
FORWARD_SPEED = 0.4                                    # скорость движения к метке и от метки в м/с
VERTICAL_SPEED = 0.25                                  # скорость коррекции высоты по метке в м/с
YAW_RATE = 0.4                                         # скорость поворота к метке в рад/с
COMMAND_INTERVAL = 0.3                                 # время действия команды скорости в секундах
SEND_PERIOD = 0.2                                      # как часто повторять команду скорости, включая нулевую
ZERO_SPEED = (0.0, 0.0, 0.0, 0.0)                      # нулевая команда: vx, vy, vz, yaw_rate


def load_coefficients(path):                          # функция для загрузки коэффициентов калибровки камеры
    cv_file = cv2.FileStorage(str(path), cv2.FILE_STORAGE_READ) # открываем файл для чтения данных калибровки

    if not cv_file.isOpened():                         # проверяем, что файл калибровки удалось открыть
        raise FileNotFoundError("Не удалось открыть файл калибровки: " + str(path)) # сообщаем, какой файл не найден

    camera_matrix = cv_file.getNode("mtx").mat()      # считываем матрицу камеры из файла
    dist_coeffs = cv_file.getNode("dist").mat()       # считываем коэффициенты искажений из файла
    cv_file.release()                                 # закрываем файл после чтения

    if camera_matrix is None or dist_coeffs is None:   # проверяем, что файл содержит нужные данные
        raise ValueError("Файл калибровки не содержит mtx или dist") # сообщаем о неверном файле калибровки

    return camera_matrix, dist_coeffs                 # возвращаем матрицу камеры и коэффициенты искажений


def setup_terminal():                                 # функция включает чтение клавиш без нажатия Enter
    if not sys.stdin.isatty():                        # проверяем, что программа запущена в терминале
        return None                                   # возвращаем None, если терминал недоступен

    settings = termios.tcgetattr(sys.stdin)           # сохраняем текущие настройки терминала
    tty.setcbreak(sys.stdin.fileno())                 # включаем режим чтения по одному символу
    return settings                                   # возвращаем настройки, чтобы восстановить их позже


def restore_terminal(settings):                       # функция восстанавливает настройки терминала
    if settings is not None:                          # проверяем, что настройки были сохранены
        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, settings) # возвращаем терминал в исходный режим


def read_terminal_key():                              # функция проверяет, нажата ли клавиша в терминале
    if select.select([sys.stdin], [], [], 0)[0]:       # проверяем, есть ли введенная строка
        return sys.stdin.read(1).lower()               # считываем один символ и приводим его к нижнему регистру
    return None                                        # возвращаем None, если клавиша не нажата


def get_marker_zone(x_center, frame_width):            # функция определяет, в какой зоне кадра находится метка
    if x_center < frame_width / 3:                     # проверяем, что центр метки слева от центральной зоны
        return "left"                                  # возвращаем левую зону
    if x_center > frame_width * 2 / 3:                 # проверяем, что центр метки справа от центральной зоны
        return "right"                                 # возвращаем правую зону
    return "center"                                    # возвращаем центральную зону


def get_vertical_zone(y_center, frame_height):          # функция определяет, выше или ниже центра находится метка
    if y_center < frame_height / 3:                    # проверяем, что центр метки выше центральной зоны
        return "top"                                   # возвращаем верхнюю зону
    if y_center > frame_height * 2 / 3:                # проверяем, что центр метки ниже центральной зоны
        return "bottom"                                # возвращаем нижнюю зону
    return "center"                                    # возвращаем центральную зону


def setup_servo_camera(angle):                         # функция устанавливает угол сервокамеры перед полетом
    if ServoCamera is None:                            # проверяем, доступно ли управление сервокамерой
        print("Сервокамера недоступна в текущей конфигурации SDK") # выводим предупреждение
        return None                                    # возвращаем None вместо объекта сервокамеры

    try:                                               # пробуем подключиться к сервокамере
        servo_camera = ServoCamera()                   # создаем объект сервокамеры

        if ServoPriority is not None:                  # проверяем, доступен ли класс приоритета сервокамеры
            success = servo_camera.set_angle(angle, ServoPriority.MEDIUM) # устанавливаем угол с обычным приоритетом
        else:                                          # если класс приоритета недоступен
            success = servo_camera.set_angle(angle)    # устанавливаем угол без указания приоритета

        if success:                                    # проверяем, удалось ли отправить команду
            print(f"Камера установлена на {angle}°")   # выводим установленный угол
        else:                                          # если команда не была отправлена
            print(f"Не удалось установить камеру на {angle}°") # выводим предупреждение

        return servo_camera                            # возвращаем объект сервокамеры
    except Exception as error:                         # если сервокамера недоступна
        print("Сервокамера недоступна:", error)        # выводим причину в терминал
        return None                                    # возвращаем None вместо объекта сервокамеры


class VideoProcessingThread(threading.Thread):        # класс для обработки видео в отдельном потоке
    def __init__(self, camera_matrix, dist_coeffs):   # метод инициализации класса
        super().__init__()                            # вызываем инициализацию родительского класса Thread
        self.camera_matrix = camera_matrix            # сохраняем матрицу камеры
        self.dist_coeffs = dist_coeffs                # сохраняем коэффициенты искажений
        self.running = True                           # флаг работы потока
        self.latest_coordinates = None                # переменная для хранения координат найденной метки
        self.x_center = None                          # переменная для хранения центра метки по оси X
        self.y_center = None                          # переменная для хранения центра метки по оси Y
        self.frame_width = None                       # переменная для хранения ширины кадра
        self.frame_height = None                      # переменная для хранения высоты кадра
        self.marker_ids = None                        # переменная для хранения ID найденных меток
        self.debug_in_sky = False                     # переменная показывает, находится ли дрон в воздухе
        self.debug_distance = None                    # переменная хранит расчетную дистанцию до метки
        self.debug_zone = "lost"                      # переменная хранит зону кадра с найденной меткой
        self.debug_vertical_zone = "lost"             # переменная хранит вертикальную зону кадра с найденной меткой
        self.debug_vy = 0                             # переменная хранит команду скорости вперед/назад
        self.debug_vz = 0                             # переменная хранит команду скорости вверх/вниз
        self.debug_yaw_rate = 0                       # переменная хранит команду скорости поворота
        self.debug_status = "waiting takeoff"         # переменная хранит текстовое состояние алгоритма

        self.aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_ARUCO_ORIGINAL) # выбираем словарь ArUco-меток
        self.aruco_params = cv2.aruco.DetectorParameters()                         # создаем параметры для поиска ArUco-меток
        self.aruco_detector = cv2.aruco.ArucoDetector(self.aruco_dict, self.aruco_params) # создаем детектор ArUco-меток

        self.size_of_marker = 0.1                   # задаем размер стороны ArUco-метки в метрах
        self.points_of_marker = np.array([          # задаем координаты углов ArUco-метки
            (self.size_of_marker / 2, -self.size_of_marker / 2, 0),
            (-self.size_of_marker / 2, -self.size_of_marker / 2, 0),
            (-self.size_of_marker / 2, self.size_of_marker / 2, 0),
            (self.size_of_marker / 2, self.size_of_marker / 2, 0)
        ], dtype=np.float32)                         # используем тип float32, который подходит для функций OpenCV

        self.camera = Camera()                      # создаем экземпляр класса Camera для получения кадров с камеры
        self.viewer = ImageViewer()                 # создаем экземпляр класса ImageViewer для трансляции изображения

    def update_debug(self, in_sky, distance, zone, vertical_zone, vy, vz, yaw_rate, status): # метод обновляет диагностические данные для видео
        self.debug_in_sky = in_sky                  # сохраняем состояние полета
        self.debug_distance = distance              # сохраняем дистанцию до метки
        self.debug_zone = zone                      # сохраняем зону кадра
        self.debug_vertical_zone = vertical_zone    # сохраняем вертикальную зону кадра
        self.debug_vy = vy                          # сохраняем команду скорости вперед/назад
        self.debug_vz = vz                          # сохраняем команду скорости вверх/вниз
        self.debug_yaw_rate = yaw_rate              # сохраняем команду поворота
        self.debug_status = status                  # сохраняем описание текущего действия

    def draw_debug_overlay(self, frame):            # метод рисует диагностическую информацию на кадре
        height, width = frame.shape[:2]             # получаем высоту и ширину кадра
        cv2.line(frame, (width // 3, 0), (width // 3, height), (255, 255, 0), 1) # рисуем левую границу центральной зоны
        cv2.line(frame, (width * 2 // 3, 0), (width * 2 // 3, height), (255, 255, 0), 1) # рисуем правую границу центральной зоны
        cv2.line(frame, (width // 2, 0), (width // 2, height), (80, 80, 80), 1) # рисуем центр кадра по оси X
        cv2.line(frame, (0, height // 3), (width, height // 3), (255, 255, 0), 1) # рисуем верхнюю границу центральной зоны
        cv2.line(frame, (0, height * 2 // 3), (width, height * 2 // 3), (255, 255, 0), 1) # рисуем нижнюю границу центральной зоны
        cv2.line(frame, (0, height // 2), (width, height // 2), (80, 80, 80), 1) # рисуем центр кадра по оси Y

        if self.x_center is not None and self.y_center is not None: # проверяем, найдена ли метка на кадре
            cv2.arrowedLine(frame, (width // 2, height // 2), (self.x_center, self.y_center), (0, 255, 255), 2) # показываем смещение метки от центра кадра

        marker_text = "marker: lost"               # создаем строку состояния метки
        if self.marker_ids is not None:             # проверяем, есть ли найденные ID
            marker_text = f"marker: id={self.marker_ids}" # добавляем ID найденных меток

        position_text = "pos: no pose"              # создаем строку координат метки
        if self.latest_coordinates is not None:      # проверяем, рассчитаны ли координаты метки
            x, y, z = self.latest_coordinates        # получаем координаты метки
            position_text = f"pos: x={x:.2f} y={y:.2f} z={z:.2f}" # формируем строку координат

        distance_text = "dist: n/a"                 # создаем строку дистанции
        if self.debug_distance is not None:          # проверяем, рассчитана ли дистанция
            distance_text = f"dist: {self.debug_distance:.2f} m" # формируем строку дистанции

        flight_text = "flight: in sky" if self.debug_in_sky else "flight: waiting" # формируем строку состояния полета
        center_text = f"center: x={self.x_center} y={self.y_center} zone={self.debug_zone}/{self.debug_vertical_zone}" # формируем строку положения метки в кадре
        command_text = f"cmd: vy={self.debug_vy:.2f} vz={self.debug_vz:.2f} yaw={self.debug_yaw_rate:.2f}" # формируем строку команды движения

        lines = [marker_text, position_text, distance_text, center_text, command_text, flight_text, self.debug_status] # собираем строки вывода
        y = 24                                      # задаем начальную координату текста

        for line in lines:                         # перебираем строки диагностики
            cv2.putText(frame, line, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 0), 3) # рисуем темную обводку текста
            cv2.putText(frame, line, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1) # рисуем светлый текст
            y += 22                                # переходим к следующей строке

    def run(self):                                  # метод, который выполняется после запуска потока
        while self.running:                         # запускаем цикл, пока поток должен работать
            try:                                    # пробуем получить кадр с камеры
                frame = self.camera.get_cv_frame(timeout=0.1) # получаем один кадр с камеры
                                                             # timeout=0.1 - короткое ожидание, чтобы поток быстрее останавливался
            except TimeoutError:                    # если кадр не успел прийти
                continue                            # пропускаем текущую итерацию цикла

            if frame is None:                       # проверяем, что кадр не был получен
                continue                            # пропускаем текущую итерацию цикла

            self.frame_width = frame.shape[1]        # сохраняем ширину кадра
            self.frame_height = frame.shape[0]       # сохраняем высоту кадра
            corners, ids, rejected = self.aruco_detector.detectMarkers(frame) # ищем ArUco-метки на изображении

            if ids is not None and len(corners) > 0: # проверяем, что ArUco-метка найдена
                x_center = int(sum(point[0] for point in corners[0][0]) / 4) # вычисляем центр метки по оси X
                y_center = int(sum(point[1] for point in corners[0][0]) / 4) # вычисляем центр метки по оси Y
                self.marker_ids = ids.flatten().tolist() # сохраняем ID найденных меток

                cv2.circle(frame, (x_center, y_center), 5, (0, 0, 255), -1) # рисуем центр найденной метки
                cv2.aruco.drawDetectedMarkers(frame, corners, ids)          # рисуем найденную ArUco-метку

                image_points = corners[0].reshape(-1, 2) # получаем координаты углов первой найденной метки

                success, rvecs, tvecs = cv2.solvePnP( # считаем положение ArUco-метки относительно камеры
                    self.points_of_marker,
                    image_points,
                    self.camera_matrix,
                    self.dist_coeffs
                )

                if success:                         # проверяем, что положение метки успешно рассчитано
                    self.latest_coordinates = [tvecs.item(0), tvecs.item(1), tvecs.item(2)] # сохраняем координаты метки
                    self.x_center = x_center        # сохраняем центр метки по оси X
                    self.y_center = y_center        # сохраняем центр метки по оси Y
                    self.frame_width = frame.shape[1] # сохраняем ширину кадра
                    self.frame_height = frame.shape[0] # сохраняем высоту кадра
                    cv2.drawFrameAxes(frame, self.camera_matrix, self.dist_coeffs, rvecs, tvecs, 0.1) # рисуем оси метки
                else:                               # если положение метки не удалось рассчитать
                    self.latest_coordinates = None  # очищаем координаты метки
                    self.x_center = None            # очищаем центр метки по оси X
                    self.y_center = None            # очищаем центр метки по оси Y
            else:                                   # если ArUco-метка не найдена
                self.latest_coordinates = None      # очищаем координаты метки
                self.x_center = None                # очищаем центр метки по оси X
                self.y_center = None                # очищаем центр метки по оси Y
                self.marker_ids = None              # очищаем ID найденных меток

            self.draw_debug_overlay(frame)           # рисуем диагностическую информацию на кадре
            self.viewer.imshow("aruco_flight", frame, fps=30) # запускаем трансляцию изображения

    def stop(self):                                  # метод для остановки потока обработки видео
        self.running = False                        # меняем флаг работы потока
        self.camera.stop()                          # останавливаем получение кадров с камеры
        self.viewer.close()                         # останавливаем видеопоток


def wait_for_point(drone):                          # функция ожидания прилета дрона в точку
    while not drone.point_reached():                 # ждем, пока дрон не достигнет заданной точки
        time.sleep(0.1)                              # ставим небольшую паузу, чтобы не нагружать программу


setup_servo_camera(CAMERA_HORIZONTAL_ANGLE)          # устанавливаем камеру горизонтально перед запуском видео
camera_matrix, dist_coeffs = load_coefficients(DATA_PATH) # загружаем коэффициенты калибровки камеры
video_thread = VideoProcessingThread(camera_matrix, dist_coeffs) # создаем поток обработки видео
video_thread.start()                                       # запускаем поток обработки видео

drone = Pioneer()                                   # создаем экземпляр класса Pioneer, устанавливаем соединение
terminal_settings = None                            # переменная хранит исходные настройки терминала
last_debug_print_time = 0                           # переменная хранит время последнего диагностического вывода
last_speed_send_time = 0.0                           # переменная хранит время последней отправки скорости
last_sent_speed = None                               # переменная хранит последнюю отправленную скорость

try:                                                # основной код находится внутри блока try
    terminal_settings = setup_terminal()            # включаем чтение клавиш без нажатия Enter
    print("Нажмите 's' для взлета")                 # выводим инструкцию для пользователя
    print("Нажмите 'q' для посадки и завершения")   # выводим инструкцию для пользователя

    while True:                                     # запускаем бесконечный цикл
        coordinates = video_thread.latest_coordinates # получаем последние координаты найденной метки
        x_center = video_thread.x_center            # получаем центр найденной метки по оси X
        y_center = video_thread.y_center            # получаем центр найденной метки по оси Y
        frame_width = video_thread.frame_width      # получаем ширину кадра
        frame_height = video_thread.frame_height    # получаем высоту кадра
        marker_ids = video_thread.marker_ids        # получаем ID найденных меток
        key = read_terminal_key()                   # считываем нажатую клавишу из терминала
        state = drone.get_fly_state().name          # состояние дрона: ON_LAND, ARMED или IN_SKY
        in_sky = state == "IN_SKY"                  # True, если дрон уже в воздухе

        vy = 0                                      # задаем скорость вперед/назад относительно корпуса дрона
        vz = 0                                      # задаем скорость вверх/вниз
        yaw_rate = 0                                # задаем скорость поворота по курсу
        distance = None                             # создаем переменную для дистанции до метки
        zone = "lost"                               # создаем переменную для зоны кадра
        vertical_zone = "lost"                      # создаем переменную для вертикальной зоны кадра
        status = "waiting takeoff"                  # создаем текстовое состояние алгоритма

        if key == "q":                              # проверяем, нажата ли клавиша выхода
            break                                   # выходим из цикла

        if key == "s" and not in_sky:                # проверяем, нажата ли клавиша взлета
            if state == "ON_LAND":                   # если двигатели еще не включены
                drone.arm()                          # включаем двигатели
            if drone.get_fly_state().name == "ARMED": # проверяем, что двигатели включились
                drone.takeoff()                      # производим взлет
                drone.go_to_local_point(x=0, y=0, z=1.5, yaw=0, time=3) # набираем заданную высоту полета
                wait_for_point(drone)                # ждем, пока дрон поднимется на заданную высоту
                last_sent_speed = None               # после взлета сразу отправим первую команду скорости
                last_speed_send_time = 0.0           # сбрасываем таймер отправки скорости

        in_sky = drone.get_fly_state().name == "IN_SKY" # обновляем состояние после возможного взлета

        if in_sky and coordinates is not None and x_center is not None and y_center is not None and frame_width is not None and frame_height is not None: # проверяем, что дрон в воздухе и метка найдена
            distance = float(np.linalg.norm(coordinates)) # вычисляем расстояние до ArUco-метки
            zone = get_marker_zone(x_center, frame_width) # определяем зону кадра с найденной меткой
            vertical_zone = get_vertical_zone(y_center, frame_height) # определяем вертикальную зону кадра с найденной меткой
            status_parts = []                       # создаем список действий, выбранных алгоритмом

            if distance > 1.5:                      # проверяем, что дрон далеко от метки
                vy = FORWARD_SPEED                  # задаем скорость движения вперед
                status_parts.append("forward")      # сохраняем действие по дистанции
            elif distance < 1.0:                    # проверяем, что дрон слишком близко к метке
                vy = -FORWARD_SPEED                 # задаем скорость движения назад
                status_parts.append("backward")     # сохраняем действие по дистанции
            else:                                    # если дистанция находится в рабочем диапазоне
                status_parts.append("hold distance") # сохраняем действие по дистанции

            if x_center < frame_width / 3:          # проверяем, что метка находится слева на изображении
                yaw_rate = YAW_RATE                 # задаем поворот влево
                status_parts.append("turn left")    # сохраняем действие по курсу
            elif x_center > frame_width * 2 / 3:    # проверяем, что метка находится справа на изображении
                yaw_rate = -YAW_RATE                # задаем поворот вправо
                status_parts.append("turn right")   # сохраняем действие по курсу
            else:                                    # если метка находится в центральной зоне
                status_parts.append("hold yaw")     # сохраняем действие по курсу

            if y_center < frame_height / 3:          # проверяем, что метка находится выше центральной зоны кадра
                vz = VERTICAL_SPEED                  # задаем набор высоты
                status_parts.append("up")            # сохраняем действие по высоте
            elif y_center > frame_height * 2 / 3:    # проверяем, что метка находится ниже центральной зоны кадра
                vz = -VERTICAL_SPEED                 # задаем снижение
                status_parts.append("down")          # сохраняем действие по высоте
            else:                                    # если метка находится по центру по вертикали
                status_parts.append("hold height")   # сохраняем действие по высоте

            status = ", ".join(status_parts)         # собираем текстовое описание действия
        elif in_sky:                                 # если дрон в воздухе, но метка не найдена
            status = "marker lost, hover"           # сохраняем состояние потери метки

        video_thread.update_debug(in_sky, distance, zone, vertical_zone, vy, vz, yaw_rate, status) # обновляем диагностику для видеопотока

        if in_sky and time.time() - last_debug_print_time > DEBUG_PRINT_INTERVAL: # проверяем, пора ли печатать диагностику
            if coordinates is None or x_center is None or y_center is None or frame_width is None or frame_height is None: # проверяем, видна ли метка
                print("ArUco: метка не найдена -> команда зависания") # выводим состояние потери метки
            else:                                      # если метка найдена
                print(                                # выводим данные о метке и команде движения
                    f"ArUco: id={marker_ids}, "
                    f"dist={distance:.2f} м, "
                    f"x={x_center}/{frame_width}, y={y_center}/{frame_height}, "
                    f"zone={zone}/{vertical_zone}, "
                    f"vy={vy:.2f}, vz={vz:.2f}, yaw={yaw_rate:.2f}, "
                    f"action={status}"
                )
            last_debug_print_time = time.time()        # запоминаем время вывода

        speed_command = (0.0, vy, vz, yaw_rate)      # собираем команду скорости: vx, vy, vz, yaw_rate
        send_due = time.monotonic() - last_speed_send_time >= SEND_PERIOD # проверяем, пора ли повторить команду
        speed_changed = speed_command != last_sent_speed # новую скорость отправляем сразу

        if in_sky and (speed_changed or send_due):   # проверяем, нужно ли отправить команду скорости
            drone.set_manual_speed_body_fixed(*speed_command, COMMAND_INTERVAL) # отправляем команду скорости
            last_speed_send_time = time.monotonic()  # запоминаем время отправки команды
            last_sent_speed = speed_command          # запоминаем последнюю отправленную скорость

        time.sleep(0.05)                            # ставим небольшую паузу, чтобы не нагружать программу

finally:                                            # блок finally выполнится при завершении программы
    restore_terminal(terminal_settings)             # восстанавливаем настройки терминала
    video_thread.stop()                             # останавливаем поток обработки видео
    video_thread.join()                             # ждем завершения потока обработки видео

    if drone.get_fly_state().name == "IN_SKY":      # проверяем, находится ли дрон в воздухе
        drone.set_manual_speed_body_fixed(*ZERO_SPEED, COMMAND_INTERVAL) # отправляем нулевую скорость перед посадкой
        drone.land()                                # производим посадку

    drone.close_connection()                        # закрываем соединение с дроном