Осенью мы с друзьями участвовали в хакатоне DIGITAL SUPERHERO от организаторов хакатона ИТС и СЦ, о котором я писал в статье. Хакатон проходил с 18 по 23 сентября 2020 года. Мы выбрали задачу Распознавание аномалий (объектов и инцидентов) на фотоматериалах, полученных с беспилотных летательных аппаратов (БПЛА) в треке "Разработка алгоритмов распознавания. В рамках кейса нужно было разработать алгоритм по автоматическому выявлению аномалий на изображении и разработать веб-интерфейс для загрузки и разметки изображений. В этой статье я бы хотел рассказать о нашем решении для детекции аномалий с помощью модели YOLOv5 в виде практического туториала. Кому интересно, прошу под кат.

Описание задачи

Задачу предоставила нефтяная компания Татнефть. Татнефть использует БПЛА для мониторинга районов разработки нефти. В таких районах могут возникнуть различные аномалии такие, как разлив нефти, вскрытие грунта, складирование материала, работа спец. техники. Обычно детекция таких аномалий выполняется операторами вручную, что требует много времени. Такую задачу имеет смысл автоматизировать с применением методов компьютерного зрения. Наша задача была представить такой алгоритм автоматического выявления аномалий на изображении с БПЛА.

Исследование данных

К сожалению, мы не смогли найти датасет изображений с БПЛА, на котором имеются подобные аномалии. Поэтому пришлось использовать датасет, предоставленный организаторами.

Нам был предоставлен довольно маленький набор данных, сотоящий из изображений размера 6000x4000, разделенных на 4 вышеупомянутых класса аномалий. Классы оказались сильно несбалансированными: например много семплов пришлось на класс складирование материала (15 изображений и до 5 объектов класса на изображении) и очень мало на класс работы спец. техники (6 изображений). Также было много изображений, не включающих ни одной аномалии (негативные семплы). Стоит отметить, что на одном изображении встречаются только аномалии одного класса, что сильно упрощает задачу.

Пример изображения

Подготовка данных для обучения

Как я уже упомянул нам предоставили изображения размера 6000x4000 со следующим распределением по классам: складирование материала - 15 семплов, вскрытие грунта - 8, разлив нефти - 15, работа спец. техники - 6 и 121 изображение без аномалий.

Для обучения моделей детекции объектов необходимо тысячи и даже миллионы образцов. Поэтому мы решили резать все изображения на тайлы размером 600x400 с перекрытием 25% методом скользящего окна.

Пример тайла

Для разметки мы использовали онлайн сервис makesense.ai. Его интерфейс очень прост и не требуется регистрации. Просто нажимаем кнопку Get started и загружаем файлы папкой через драг-н-дроп или файд аплоадер, создаем список меток и нажимаем Start project. Также можно выбрать автоматическую разметку изображений на основе искусственного интеллекта. После окончания разметки результат можно экспортировать в формате YOLO.

Изначально у нас было 27K изображений (после нарезки на тайлы). Обучение модели YOLOv5 на таком большом наборе данных занимает очень много времени, поэтому мы уменьшили датасет до 10K.

Затем мы применили аугментацию данных с помощью библиотеки imgaug для увеличения датасета и улучшения качества детекции. О важности аугментации данных в задаче детекции можно прочитать здесь. YOLOv5 применяет некоторые техники аугментации "из коробки" (например mosaic data augmentation). Мы в это задаче использовали поворот, горизонтальный и вертикальное отражение (flip), сдвиг по горизонтали и вертикали. Наш скрипт для аугментации данных доступен здесь. Для аугментации данных также можно использовать сервис roboflow.

Примеры аугментации

Выбор модели

Здесь мы выбрали YOLOv5, недавний релиз серии YOLO. По информации из сети YOLO обходит по точности другие SOTA модели на датасете COCO и представляет хороший баланс между точностью и скоростью выполения. Здесь представлен хороший сравнительный бенчмарк различных SOTA моделей.

Обучение модели

Мы использовали PyTorch реализацию YOLOv5 с этого github репозитория. Подробную инструкцию по обучению модели на своем датасете можно найти здесь. Для обучения лучше использовать Google Colab, поскольку он предоставляет бесплатно GPU (правда есть ограничение на 8 часов в сутки).

Просто клонируем репозиторий и устанавливаем зависимости:

pip install -qr requirements.txt

Скопируем файл coco128.yaml в папке data чтобы в дальнейшем

cd data && cp coco128.yaml anomalies.yaml

и начнем его изменять под нашу задачу. Сначала зададим пути к обучающему и валидационному датасетам

train: ../dataaug/images/train/  val: ../dataaug/images/val/

Изменим число классов для обучения

# number of classes: 4

и укажем имена наших классов

# class names: ['excavation', 'service', 'spill', 'storage']

Также изменим число классов в файле models/yolov5s.yaml

nc: 4 # number of classes

И запусим обучение

python train.py --img 416 --batch 64 --epochs 20 --data ./data/anomalies.yaml --cfg ./models/yolov5s.yaml --weights ''

Здесь мы используем модель yolov5s.

Модель yolov5 создает лог обучения в папке runs

cd yolov5 && ls -l runs 
drwx------ 2 root root 4096 Sep 21 09:52 exp0
drwx------ 2 root root 4096 Sep 21 15:07 exp1
…

Метрики сохраняются в виде изображения results.png

Распределение меток классов в данных сохраняются в виде изображения labels.png

Примеры одного из самых ранних результатов детекции на нашем датасете с дефолтными значениями параметров обучения YOLOv5

Оптимизация производительности

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

На этом все. Удачи в обучении YOLOv5 для детекции объектов и до новых встреч.