Простое сопровождение человека

import cv2                                             # библиотека cv2 нужна для обработки изображений
import time                                            # библиотека time содержит функции для работы со временем

from pioneer_sdk2 import Pioneer, Camera, ImageViewer, CameraType # импортируем классы из библиотеки pioneer_sdk2
from pioneer_rknn import YoloPose                      # импортируем модель YOLO Pose из библиотеки pioneer_rknn

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


MODEL_NAME = "yolov8n-pose"                            # имя модели YOLO Pose в реестре моделей Pioneer-RKNN
IMG_SIZE = (640, 640)                                  # размер изображения, который подается на вход нейросети

CAMERA_TOP_ANGLE = 25                                  # верхнее положение сервокамеры в градусах
FLIGHT_HEIGHT = 1.5                                    # рабочая высота полета в метрах
TAKEOFF_SETTLE_TIME = 0.5                              # пауза после выхода на рабочую высоту

COMMAND_INTERVAL = 0.3                                 # время действия команды скорости в секундах
SEND_PERIOD = 0.2                                      # как часто повторять команду скорости, включая нулевую
ZERO_SPEED = (0.0, 0.0, 0.0, 0.0)                      # нулевая команда: vx, vy, vz, yaw_rate

MAX_YAW_RATE = 0.5                                     # максимальная скорость поворота по курсу в рад/с

YAW_KP = 0.003                                         # коэффициент поворота к центру кадра

last_speed_send_time = 0.0                             # время последней отправки скорости
last_sent_speed = None                                 # последняя отправленная скорость


def limit(value, min_value, max_value):                # функция ограничивает значение заданным диапазоном
    return max(min_value, min(value, max_value))        # возвращаем значение внутри диапазона min_value..max_value


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

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

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

        if success:
            print(f"Камера установлена на {CAMERA_TOP_ANGLE}°")
        else:
            print(f"Не удалось установить камеру на {CAMERA_TOP_ANGLE}°")
    except Exception as error:
        print("Сервокамера недоступна:", error)


def resize_for_model(frame):                           # функция подготавливает кадр для нейросети
    frame = cv2.resize(frame, IMG_SIZE)                 # изменяем размер кадра до 640x640
    return frame.reshape(1, IMG_SIZE[1], IMG_SIZE[0], 3) # добавляем размерность batch для RKNN


def get_person_boxes(frame, detections):                # функция возвращает рамки найденных людей
    if detections is None:                              # проверяем, что модель вернула результат
        return []                                       # возвращаем пустой список, если результата нет

    frame_height, frame_width = frame.shape[:2]         # получаем размеры исходного кадра
    scale_x = frame_width / IMG_SIZE[0]                 # масштаб по ширине
    scale_y = frame_height / IMG_SIZE[1]                # масштаб по высоте
    boxes = []                                          # список рамок найденных людей

    for detection in detections:                        # перебираем найденные объекты
        x1 = detection.xmin * scale_x                   # левая граница рамки
        y1 = detection.ymin * scale_y                   # верхняя граница рамки
        x2 = detection.xmax * scale_x                   # правая граница рамки
        y2 = detection.ymax * scale_y                   # нижняя граница рамки

        if x2 > x1 and y2 > y1:                         # проверяем, что рамка корректная
            boxes.append((x1, y1, x2, y2))              # сохраняем рамку

    return boxes                                        # возвращаем список рамок


def select_main_box(boxes):                             # функция выбирает самого крупного человека
    if not boxes:                                       # проверяем, есть ли найденные рамки
        return None                                     # возвращаем None, если человека нет

    def box_area(box):                                  # функция считает площадь рамки
        x1, y1, x2, y2 = box                            # получаем координаты рамки
        return (x2 - x1) * (y2 - y1)                    # возвращаем площадь рамки

    return max(boxes, key=box_area)                     # выбираем рамку с максимальной площадью


def get_tracking_speed(box, frame_width):               # функция считает скорость для сопровождения человека
    x1, _, x2, _ = box                                  # получаем горизонтальные координаты рамки человека
    box_center_x = (x1 + x2) / 2                        # центр рамки по X

    yaw_error = frame_width / 2 - box_center_x          # ошибка по горизонтали: человек левее/правее центра

    vx = 0.0                                            # по оси X корпуса в этом примере не двигаемся
    vy = 0.0                                            # вперед/назад не летим: расстояние по размеру рамки ненадежно
    vz = 0.0                                            # высоту удерживает автопилот после выхода на рабочую высоту
    yaw_rate = limit(YAW_KP * yaw_error, -MAX_YAW_RATE, MAX_YAW_RATE) # скорость поворота

    return vx, vy, vz, yaw_rate                         # возвращаем команду скорости


def send_speed_command(drone, speed, immediate=True, force=False): # функция отправляет команду скорости
    global last_speed_send_time, last_sent_speed         # используем общие переменные последней команды

    send_due = time.monotonic() - last_speed_send_time >= SEND_PERIOD # пора повторить команду
    speed_changed = speed != last_sent_speed             # новая команда отличается от предыдущей

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


def draw_box(frame, box, speed):                         # функция рисует рамку и текущую команду
    x1, y1, x2, y2 = [int(value) for value in box]        # переводим координаты рамки в целые числа
    _, vy, vz, yaw_rate = speed                          # берем скорости для вывода на кадр

    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) # рисуем рамку человека
    cv2.circle(frame, ((x1 + x2) // 2, (y1 + y2) // 2), 4, (0, 255, 255), -1) # рисуем центр рамки
    cv2.putText(frame, f"vy={vy:.2f} vz={vz:.2f} yaw={yaw_rate:.2f}", (20, 40),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2) # выводим команду скорости


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


def start_flight(drone):                                # функция запускает дрон и поднимает его на рабочую высоту
    if not drone.arm():                                 # включаем двигатели
        raise RuntimeError("Не удалось включить двигатели")

    if not drone.takeoff():                             # выполняем взлет
        raise RuntimeError("Не удалось взлететь")

    drone.go_to_local_point(x=0, y=0, z=FLIGHT_HEIGHT, yaw=0, time=3) # набираем рабочую высоту
    wait_for_point(drone)                                # ждем, пока дрон поднимется на заданную высоту
    time.sleep(TAKEOFF_SETTLE_TIME)                     # ждем стабилизации после взлета


def main():                                             # главная функция программы
    global last_speed_send_time, last_sent_speed         # используем переменные отправки скорости

    drone = None                                        # объект управления дроном
    camera = None                                       # объект камеры
    viewer = None                                       # объект RTSP-трансляции
    model = None                                        # объект нейросети

    try:
        setup_servo_camera()                            # устанавливаем камеру в верхнее положение
        camera = Camera(camera_type=CameraType.MAIN)    # создаем объект основной камеры
        viewer = ImageViewer()                          # создаем объект видеопотока
        model = YoloPose(model_name=MODEL_NAME)         # загружаем модель YOLO Pose
        drone = Pioneer()                               # подключаемся к дрону

        print("Запущено сопровождение человека")
        print("Остановить программу можно сочетанием Ctrl+C")

        start_flight(drone)                             # выполняем взлет
        last_sent_speed = None                          # после взлета первая команда должна отправиться сразу
        last_speed_send_time = 0.0                      # сбрасываем таймер отправки скорости

        while True:
            try:
                frame = camera.get_cv_frame(timeout=0.05) # получаем кадр с камеры
            except TimeoutError:
                send_speed_command(drone, ZERO_SPEED)   # если кадра нет, зависаем
                continue

            if frame is None:                           # если кадр не получен
                send_speed_command(drone, ZERO_SPEED)   # отправляем нулевую скорость
                continue

            model_input = resize_for_model(frame)       # подготавливаем кадр для нейросети
            detections = model.run([model_input])       # запускаем распознавание человека
            boxes = get_person_boxes(frame, detections) # получаем рамки найденных людей
            main_box = select_main_box(boxes)           # выбираем самого крупного человека

            if main_box is None:                        # если человек не найден
                speed = ZERO_SPEED                      # задаем нулевую скорость
                send_speed_command(drone, speed)        # отправляем команду зависания
                cv2.putText(frame, "person not found", (20, 40),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
            else:
                speed = get_tracking_speed(main_box, frame.shape[1]) # считаем скорость
                send_speed_command(drone, speed, immediate=False) # отправляем скорость по таймеру
                draw_box(frame, main_box, speed)         # рисуем рамку и команду

            viewer.imshow(name="human_tracking", frame=frame, fps=30) # публикуем кадр

    except KeyboardInterrupt:
        print("Остановка программы, производится посадка")

    except Exception as error:
        print("Ошибка:", error)

    finally:
        if drone is not None:                           # если соединение с дроном было создано
            state = drone.get_fly_state().name          # читаем состояние дрона
            if state == "IN_SKY":                       # если дрон в воздухе
                send_speed_command(drone, ZERO_SPEED, force=True) # останавливаем движение
                drone.land()                            # выполняем посадку
            elif state == "ARMED":                      # если двигатели включены, но взлета не было
                drone.disarm()                          # выключаем двигатели

        if viewer is not None:                          # если видеопоток был создан
            viewer.close()                              # закрываем RTSP-трансляцию
        if camera is not None:                          # если камера была создана
            camera.stop()                               # останавливаем камеру
        if drone is not None:                           # если соединение было создано
            drone.close_connection()                    # закрываем соединение с дроном
        if model is not None:                           # если модель была создана
            model.release()                             # освобождаем ресурсы модели


if __name__ == "__main__":                             # проверяем, что файл запущен как программа
    main()                                             # запускаем главную функцию