Если вам понадобилось обучить нейронную сеть обнаруживать объекты на изображении или в видеопотоке, то для вас есть хорошие новости - уже существуют простые в использовании готовые инструменты. Наверняка кто-то из читателей такими пользуется. Например модель YOLOv8 хорошо тренируется на новых данных "из коробки". И, наверное, права была посетитель выставки, которая около стенда демонстрирующего обнаружение животных на видео c БПЛА сказала: "Что тут такого. Я за 10 строчек кода такое сделаю".
Но что, если сделал все по инструкции, а оно не работает так как хочется?

В статье покажу, какие вопросы могут возникнуть в процессе обучения нейронной сети для решения задачи обнаружения людей на снимках с дрона. Здесь не будет глубокой математики, но предполагается, что у читателя есть общее представление, что такое object detection.

Задача обнаружения людей и других объектов на снимках с беспилотных летательных аппаратов не новая. На эту тему регулярно проводятся конкурсы и написано много статей (например Поиск пропавших людей на снимках лесного массива, полученных с помощью БПЛА или ещё один разбор задачи Цифрового Прорыва). Поэтому правильным будет сначала изучить существующие подходы и решения, перед тем как приступить к разработке своего. Мы же не будем заниматься таким анализом (хотя он сделан и потянет на отдельную статью), т.к. данная статья больше про применение готовых фреймворков для решения практической задачи за "10 строк кода".

Приступим.

Подготовка датасета

Любое обучение нейронной сети начинается с подготовки данных для обучения. В случае object detection это датасет с файлами изображений и размеченными на них объектами.

Исчерпывающий список способов получить датасет изображений:

• Найти готовый датасет в свободном доступе или купить.
  Большой выбор бесплатных наборов данных для компьютерного зрения можно найти здесь https://universe.roboflow.com/. Недостаток - много повторов и некачественной разметки.

  

• Заказать сбор датасета в открытом доступе под ваши требования.
  От ~5 руб изображение. Не буду рекламировать, но если надо напишите в личку, посоветую команду, которая занимается сбором и разметкой датасетов профессионально и у кого уже заказывали. Или заказать в яндекс толока.

• Собрать самому через поисковики или самостоятельно снимать.

• Синтезировать изображения (особенно применимо для редких объектов).
  Изображения можно собрать в играх с хорошей графикой, создать 3D сцены в blender или в движках Unity, Unigine.

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

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

Мы его разбили на тренировочную, валидационную и тестовую выборку. Тренировочная будет использоваться непосредственно для обучения нейронной сети. Валидационная для проверки, что модель именно научилась выделять признаки, а не выучила "наизусть" тренировочные данные. Это называется проблемой "переобучения". Возникает обычно при малом размере или низком качестве датасета. Под низким качеством имею ввиду не качество разметки, а например если датасет состоит в основном из соседних кадров видео, снятого при одинаковом освещении, погоде и т.п. Нейронная сеть, как и человек, всегда пытается пойти по более простому пути. И если "выучить" датасет проще, чем выделить в объектах обобщающие признаки, то нейронная сеть этим воспользуется. Аналогия - зубрить предмет вместо того, чтобы понять. Тестовая выборка - для финальной проверки качества обучения. Включает данные, которые вообще не использовались во время обучения. Важно следить, чтобы изображения не пересекались между выборками.

Выбор модели

Мы будем использовать модели YOLOv5 и YOLOv8, потому что они бесплатные, хорошо показывают себя в сравнительных тестах и популярные. Обе YOLO от одного разработчика Ultralytics, реализованы на современной библиотеке глубокого обучения PyTorch. YOLOv8 позиционируется как более новая и современная по сравнению с 5й версией, хотя развиваются активно обе. Можно было ограничиться и одной моделью, мы же выберем ту, которая в процессе обучения покажет лучше результат в нашей задаче.

Разметка YOLO представляет из себя текстовые файлы с информацией о классе объекта и координатах рамки:

<номер класс> <x центра> <y центра> <ширина> <высота>

Координаты указываются относительные, от 0 до 1, чтобы не зависеть от размера изображения.

Например для image1.jpg должен быть соответствующий image1.txt, в котором может быть следующая информация:

1 0.716797 0.395833 0.216406 0.147222 
0 0.687109 0.379167 0.255469 0.158333 
1 0.420312 0.395833 0.140625 0.166667

Тренировка модели

Начнем с тренировки YOLOv5. Процесс тренировки на своих данных просто и понятно описан на сайте разработчика. Используем готовый пример запуска тренировки из документации. Только укажем количество эпох 30 и ссылку на наш файл конфигурации, в котором содержатся пути к изображениям и файлам разметки датасета:

python train.py --img 640 --epochs 30 --data data.yaml --weights yolov5s.pt

Ждем, и получаем результат тренировки: mAP50 = 5.66e-05. Метрика mAP (mean average precision) используется для оценки качества обученной модели обнаружения объектов и показывает, на сколько безошибочно нейронная сеть смогла обнаружить объекты на изображениях из валидационной выборки. Есть еще метрика "точность" (Precision) и другие, но она зависит от выбранного порога уверенности. Например при пороге 95% сеть найдет мало объектов с такой уверенностью, а при пороге 5% будет много ложных срабатываний. И если сказать, моя модель обнаружила 80% объектов, возникает вопрос "с каким порогом уверенности?". Метрика mAP более универсальна, т.к. является средним значением изменения точности в диапазоне разных порогов. Для простоты: mAP=0 модель ничего не находит, mAP=1 находит все объекты. Наш результат составляет тысячные доли процента, то есть практически нулевой.

Посмотрим как обучится YOLOv8. Поставим сразу 100 эпох. Код обучения YOLOv8 немного отличается от YOLOv5:

model = YOLO('yolov8s.pt')
model.train(data="data.yaml", epochs=100, imgsz=640)

После 30 эпох получили mAP=0.00798, после завершения тренировки: mAP=0.122. Уже можно сделать вывод, что "из коробки" YOLOv8 справляется с нашей задачей заметно лучше, чем YOLOv5. Но точность 12.2% означает, что из 10 условных человек 8-9 модель пропустит. Так никуда не годится.

Если внимательно посмотреть на пример изображения из датасета в начале статьи, то очевидно, что относительный размер объектов очень мелкий, и при подаче на вход модели изображений, сжатых до размера 640х640 пикселей, объекты сжимаются до единичных пикселей.

Мы можем использовать модель с входом 1280х1280 (также никто не запрещает указать в параметрах любой размер, хоть 8k, только поведение модели, оптимизированной под определенный размер будет непредсказуемым, а скорее даже предсказуемо хуже). Либо разбить изображение на фрагменты и подавать в модель по частям.

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

Мы пойдем по пути разделения изображений на фрагменты. Разделим огромные изображения на 20 отдельных фрагментов (тайлов).

Новый датасет увеличился в 20 раз, что не могло не отразиться на общем времени тренировки. Поэтому дальше тренировать будем только YOLOv8. Запускаем обучение с новым датасетом и получаем, что обучение прервано, т.к. за 50 последних эпох результат не улучшался.
Stopping training early as no improvement observed in last 50 epochs. Best results observed at epoch 18

Полученная точность mAP=0.215. Мы улучшили результат в ~2 раза, но он все еще неприемлемый. Можно выполнить несколько повторных тренировок и выбрать лучший результат, но не стоит ожидать, что он станет кратно лучше.

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

Обучение все также останавливается из-за отсутствия прогресса, не доходя до 100 эпох. Но результат mAP=0.408. Мы улучшили точность еще почти в два раза!

Заключение

Конечно это далеко не финальный результат. Дальше можно экспериментировать с настройками аугментации, добавлять синтетику, расширять датасет другими изображениями, пробовать другие размеры тайлов, учитывать, что при разбивке на тайлы могут быть разделены сами объекты и добиться приемлемой точности для нашей тестовой выборки.

Данное обучение проводилось для отладки алгоритмов поиска объектов на карте в самописном симуляторе.

Код первой версии симулятора можно найти здесь.

Спасибо, что прочитали до конца.