Распознаем беспилотники малых размеров с помощью ИИ

Беспилотные летательные аппараты (БПЛА), или дроны, становятся всё более распространёнными в различных областях — от коммерческой доставки и мониторинга сельскохозяйственных угодий до разведывательных миссий и обеспечения безопасности.

Однако, по мере роста их числа, возрастает и сложность задачи их распознавания, идентификации и трекинга, особенно когда речь идёт о малых дронах.

Сложно выявить малые дроны, которые благодаря своим компактным размерам и высокой манёвренности могут легко уклоняться от обнаружения радиочастотными и радарными методами. Эта задача становится ещё более сложной в условиях городской застройки или на местности с плотной растительностью.

Привет! Меня зовут Кирилл Кухарев и я разработчик в компании Raft. В этой статье я расскажу про разработку системы для детекции и трекинга БПЛА небольших размеров с помощью модели ИИ.

О чем пойдет речь в статье:

• Расскажу о том, как формировался набор данных для обучения;

• Опишу, по какому принципу выбирались модели из всего многообразия архитектур CV;

• Поделюсь, как происходил процесс обучения;

• Рассмотрю, какие метрики использовались для оценки качества работы моделей;

• Затрону принцип работы алгоритма;

• Приведу примеры сложных случаев и дальнейшие планы по улучшению работы алгоритма.

Откуда взялись данные и что с ними делали?

При решении конкретной узко-направленной задачи в компьютерном зрении на первом этапе всегда возникает вопрос, на чем обучать модель? Первое, что приходит в голову - это поискать общедоступные наборы данных на Kaggle или других профильных сайтах. Так я и поступил. Получилось найти несколько общедоступных наборов данных:

1. Набор данных БПЛА, созданный Mehdi Ozel в 2019 г. и
   опубликованный на Kaggle. Этот набор данных состоит из 1359
   изображений дронов и аннотаций к ним (txt, xml). Размер набора данных
   составляет 725.28 Мбайт.

2. Набор данных Amateur Unmanned Air Vehicle Detection
   опубликованный на Mendeley Data. Этот набор данных содержит
   около 4500 изображений любительских дронов. Размер набора данных
   составляет 157 Мбайт.

3. Набор данных с конкурса «Birds-vs-Drones».
   Содержит 104426 аннотированных изображений
   беспилотников.

Во втором столбце таблицы представлена статистика по количеству беспилотников разного размера в общедоступных наборах данных. Размеры изображений 1920×1080 px.

Проанализировав таблицу можно сказать, что в публичных наборах данных содержится большое количество изображений беспилотников, размер которых больше 1000 пикселей. Следовательно, такие изображения не соответствуют требованиям для решения поставленной задачи (определение БПЛА небольших размеров). Для повышения качества предсказания модели было принято решение увеличить набор данных путём формирования собственных изображений для обучения. Для этого было отснято 13 видеороликов с беспилотником, летающим в разных локациях (городская и сельская местность), при разных погодных условиях (дождь, снег, солнце, пасмурно) и на разном удалении от камеры. Общая продолжительность отснятого материала составила 19 часов. Для записи видео использовались две камеры HIKVISION HiWatch (модель DSI200 (C) 2.8 mm, модель DSI200 (B) 4mm) с разрешением в 1920x1080 px и скоростью записи кадров в 25 фпс. В исследовании использовались три модели дронов: DJI Phantom 4 Pro, DJI Mini 2 Fly More Combo, ZALA Z-16 VTOL. Далее с помощью скрипта на Python видео были проаннотированы в формат XML в полуавтоматическом режиме и разделены на кадры. Распределение количества беспилотников на изображении по размеру в общем наборе данных представлено в столбце 3 таблицы 1. Как можно заметить количество беспилотников меньшего размера в общем наборе данных стало гораздо больше. 

Что обучали ?

Следующий этап - это выбор архитектуры модели для обучения. Прежде чем перейти к рассмотрению архитектур стоит сказать, что есть два вида моделей детекции.

1. Двухэтапные (two-stage) детекторы. Как можно заметить из названия, процесс детекции происходит в два шага. На первом шаге выделяются регионы изображения, которые представляют особый интерес для дальнейшего анализа и обработки (regions of interest, ROI). На втором шаге выбранные регионы анализируются классификатором для определения принадлежности к заданным классам и регрессором, который уточняет местоположение ограничивающих рамок.

2. Одноэтапные (one-stage) детекторы. Подход, который заключается в предсказании координат и характеристик (класс, степень уверенности) фиксированного набора ограничивающих рамок.

Поскольку стоит задача детектирования объектов в режиме реального времени, то главными аспектами является скорость и точность обработки изображений. При сравнении двухэтапных и одноэтапных детекторов разумеется в скорости выигрывают одноэтапные детекторы. После того как определились с видом моделей стоит выделить еще два критерия, на основании которых отбирались модели:

1. Оценка скорости и качества работы модели детекции на общедоступных наборах данных, таких как MS COCO, PASCAL VOC.

2. Возможность обучить модель на доступных вычислительных мощностях. (Все вычисления производились на NVIDIA DGX-1)

Учитывая вышеуказанные требования были выбраны архитектуры детекторов семейства YOLO и SSD. Поскольку данная статья не посвящена рассмотрению принципов и описанию работы YOLO (You Only Look Once) и SSD (Single Shot Detector), то я лишь графически укажу, как каждая модель обрабатывает входное изображение.

В семействе YOLO были выбраны несколько версий моделей - YOLOv5 (n, s, m), YOLOv8 (n, s, m), GELAN-C и GELAN-E (модификации YOLOv9).

Модель SSD обучалась без изменений архитектуры с использованием чекпоинтов и регуляризации для предотвращения переобучения. Для модели YOLO был изменен размер входного изображения с 640 на 1280 для лучшего захвата мелких деталей. Также добавлены дополнительные слои и увеличено повторение слоев C2f (Cross Stage Partial Networks with Bottleneck) на разных уровнях пирамиды признаков. Также был произведен тюнинг гиперпараметров модели для их оптимизации. В качестве тюнинга каждая из версий YOLO обучалась 5 эпох с разными гиперпараметрами. Для примера показан сгенерированный после тюнингования файл с наилучшими гиперпараметрами модели YOLOv8n.

# 150/150 iterations complete ✅
# Best fitness hyperparameters are printed below.

lr0: 0.00652
lrf: 0.00402
momentum: 0.90918
weight_decay: 0.00039
warmup_epochs: 3.44701
warmup_momentum: 0.70386
box: 7.82823
cls: 0.63425
dfl: 1.40253
hsv_h: 0.0169
hsv_s: 0.67409
hsv_v: 0.19965
translate: 0.08151
scale: 0.4757
fliplr: 0.56036
mosaic: 0.9791

Далее каждая модель YOLO была обучена с учетом чекпоинтов и с размером батч пакета равным 32. 

Что получили и как оценили?

После обучения моделей необходимо было оценить, насколько хорошо модели могут определять беспилотник на изображении. Для оценки эффективности работы алгоритмов использовались метрики IoU, Precision, Recall, mAP@0.5 и mAP@0.5:0.95. Сначала сравнивались обученные модели внутри одной версии, а затем между версиями. Я лишь представлю результаты сравнения лучших моделей между версиями YOLO на валидационной выборке.

Однако определить качество работы модели, основываясь только на метриках F-меры и mAP недостаточно, тем более на валидационной выборке. Необходимо было изучить качество обнаружения беспилотников разных размеров на тестовой выборке.

Можно сказать, что наилучшая точность для беспилотников меньшего размера (до 100 px) у модели GELAN-C, а наилучшая полнота у YOLOv8n. Если сравнивать одновременно и по точности, и по полноте, то лучшие результаты на тестовой выборке получились у YOLOv8n. Следовательно в качестве базовой модели для создания системы определения БПЛА в реалтайме использовалась YOLOv8n. 

И будет хорошо определять?

Разумеется нет, при рассмотрении движущихся объектов в видеозаписях с камер видеонаблюдения или в режиме реального времени иногда возникают ситуации, когда объект на соседних кадрах может детектироваться прерывисто или вовсе может быть не распознанным. Для таких случаев существуют трекеры, которые могут отслеживать как один, так и несколько объектов одновременно. Трекер включает в себя последовательную оценку положения, размера и ориентации объекта на нескольких подряд идущих кадрах. В процессе изучения трекеров вместе с моделью YOLOv8n были протестированы трекеры BotSORT, ByteTrack и DeepSORT.

Наилучших результатов удалось достичь с трекером DeepSORT. В среднем, для каждого размера беспилотника точность повысилась на 3-4 п. п. Полнота повысилась, в среднем, на 4-5 п. п. При добавлении трекера акцент делался на повышение значения полноты, особенно для небольших размеров БПЛА. Для определения дронов в режиме реального времени был реализован следующий алгоритм.

На вход системе поступают кадры с видеокамеры. Далее происходит изменение размера изображения до 1080x1080 пикселей, чтобы подать это изображение на вход модели нейронной сети. Если предсказанные значения координат рамки и класса объекта меньше порогового значения достоверности, то ограничивающая рамка не отображается. Если модель определила беспилотник на пяти подряд идущих кадрах со значением

, то запускался трекер. После запуска удаляются ограничивающие рамки детектора и отображаются координаты ограничивающей рамки трекера, а также маршрут беспилотника в видеопотоке.

При инференсе системы скорость обработки 1 кадра составляла, в среднем, 150 мс (~6 фпс). Большую часть времени занимала детекция объектов, так как скорость детекции была обратно пропорциональной количеству БПЛА на изображении. Для оптимизации системы модель YOLO была переведена в формат ONNX и адаптирована для работы на графическом процессоре. В результате скорость обработки одного изображения сократилась до 25-30 мс(~32 фпс).

Какие есть нюансы?

Не смотря на то, что в большинстве случаев система точно определяет беспилотник небольших размеров на изображении, все же существуют случаи, когда детектор теряется и не может распознать дрон. К таким ситуациям можно отнести:

1. Плохая освещенность территории

2. В кадр попала только часть беспилотника

3. Размер БПЛА слишком маленький

4. Цвет БПЛА сливается с задним фоном

Для решения проблемы детекции в условиях плохой освещённости в наборе данных присутствуют изображения с дроном в инфракрасном режиме. При условии, что видеокамера поддерживает инфракрасный режим, проблема плохой освещённости будет отсутствовать.

В случае, когда объект попадает в кадр неполностью достаточно сложно однозначно отделить является ли он беспилотником или нет. В кадре должно находиться не менее 70% объекта для однозначного определения БПЛА.

Модуль однозначно определяет БПЛА размером не меньше 7x7 пикселей с точностью 83%. Затем, со снижением размера дрона, точность детекции начинает резко снижаться. Это связано с тем, что модель не может извлечь достаточное количество признаков, которые бы соответствовали именно дрону. Решением такой проблемы является камера с оптическим зумом для увеличения размеров подозрительного объекта.

Цвет беспилотника иногда может совпадать с цветом заднего фона.

Однако перечисленные ситуации были решены с помощью трекера при определении траектории движения объекта в видеопотоке.

Выводы

В процессе разработки системы был собран набор данных с беспилотниками, в котором содержится > 400 тыс. аннотированных изображений, обучены модели детекции и трекинга объектов, решены проблемные случаи. 

Работа по улучшению системы не заканчивается. На данный момент ведется обучение моделей YOLOv10, RT-DETR, поиск более совершенных трекеров и оптимизация алгоритма работы системы.

Исходный код программы для полуавтоматической разметки видео, обучения моделей YOLO и SSD расположен тут.