Как узнать принцессу среди 500 амурских тигров с помощью vision transformers

Рассказываем, как мы, ML princesses [Napoleon IT] стали победителями кейса от Минприроды «Защита редких животных» и решили задачу по созданию сервиса, способного распознавать в дикой природе особо редкий вид хищников — амурского тигра. Десятый региональный хакатон проходил в рамках федерального проекта «Искусственный интеллект».

Проблематика

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

Подготовка к хакатону

Чтобы хорошо подготовиться к хакатону, мы нашли готовые датасеты для детекции и классификации животных.

Датасет для детекции:

1. 2750 изображений;
2. разрешение 1920x1080;
3. разметка в формате x, y, width, height;
4. 1 класс — тигры.

Датасет для классификации:

1. 3080 изображений;
2. разрешение варьируется от изображения к изображению;
3. размеченные фотографии были распределены по папкам с названиями классов;
4. 5 классов — гепарды, леопарды, пантеры, ягуары и тигры.

Постановка задачи

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

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

Технический эксперт оценивает решение по следующим критериям:

1. Запускаемость кода.
2. Обоснованность выбранного метода (описание подходов к решению, их обоснование и релевантность задаче).
3. Точность работы алгоритма (возможность оценить формальной метрикой с обоснованием выбора).
4. Адаптивность/Масштабируемость.
5. Отсутствие в решении импортного ПО и библиотек, кроме свободно распространяемого с обоснованием выбора.

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

У нас было готово: предобученные модели классификации и детекции.

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

Исходные данные

Мы получили данные от организаторов. Это были размеченные фотографии с фотоловушек, расположенных в лесу. Разметка Всего было 4 класса: “принцесса”, “тигры”, “леопарды” и “другие животные”. Фотографии были сделаны как днем, так и ночью с разрешением 2048  x  1536 пикселей. На одном снимке могло быть несколько животных одного класса. К фотографиям прилагались csv с координатами для обучения детектора в формате x1, y1, x2, y2. На фотографии ниже можем наблюдать ту самую принцессу. 

Техническое решение

Наше решение позволяет распознавать виды животных и отдельных особей с высокой точностью за счет SOTA (state of the art) подходов. Оно состоит из 4 модулей:

1. Телеграм-бот.
2. Сервис распознавания.
3. Модель детекции.
4. Модель классификации.

С тем как эти модули связаны между собой можно ознакомиться ниже:

Детектор 

Для детекции животных использовали архитектуру DetectoRS с использованием фреймворка mmdetection. По нашему опыту, она довольно хорошо дообучается и хватает 3-4 эпохи для получения высоких метрик. Обучение модели заняло около 2 часов. Изначально мы обучали детектор на нескольких классах, но потом было принято решение оставить один, потому что в дальнейшем мы выделили отдельный этап классификации.

Код обучения детектора можно посмотреть по ссылке: https://www.kaggle.com/vfomenko/tiger-detectors-mmdetection-training 

Классификатор

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

Но чтобы выделить Принцессу и в дальнейшем была возможность распознавать и других особей, мы отказались от классической реализации с использованием полносвязного слоя и категориальной перекрестной энтропии. Для этого использовали подход на основе metric learning, который хорошо себя зарекомендовал в задаче распознавания лиц. Обучали модель для извлечения вектора признаков с функцией потерь ArcFace, и в дальнейшем сравнивали извлеченные из кадра признаки с теми, которые есть в базе данных по косинусной метрике. Вектор признаков брался из первого токена, который обычно используется для классификации, вместо различных вариаций глобальных пулингов.

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

С кодом обучения классификатора можно ознакомиться по ссылке: https://www.kaggle.com/emorkrin/animal-classification.

Метрики

По детекции мы получили 0.98 mAP и 0.99 recall

Для классификации считали три метрики: precision, recall и accuracy. Получили следующие результаты:

В итоге мы находим 99% животных и принцессу в 97% случаев с точностью 96% 

Backend

Мы реализовали бэкэнд приложение с помощью фреймворка Flask, позволяющего создавать Rest API сервисы. Создали endpoint,  который на вход принимает один кадр. В response получали json с меткой класса и координатами животного.

Frontend

Для демонстрации работы сервиса было принято решение использовать телеграм-бота. Это удобный формат демонстрации работы сервиса как для разработчика с точки зрения скорости разработки, так и для конечного пользователя, потому что никакого дополнительного ПО, кроме телеграма, устанавливать не нужно. Бот будет получать данные от юзера, взаимодействовать с нашим сервисом и выдавать краткую сводку о распознавании. В бота пользователь мог загрузить фотографию или видео. Если это видео, то оно нарезается на несколько кадров с задержкой в 3 секунды и обрабатывается как обычное изображение. Такой таймлаг был выбран для того, чтобы не засорять ленту пользователю. Для пользователя это выглядело так:

Итог

Наше решение полностью соответствует заявленным критериям:

1. Окружения всех наших сервисов зафиксированы через docker.
2. Наше решение с детекцией и последующим сопоставлением животного с известными нам позволяет идентифицировать всех особей в кадре.
3. Добились метрик работы моделей выше 97%.
4. Наше решение масштабируется на любое количество особей за счет подхода metric learning.
5. В решении были использованы только open source технологии.

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

С кодом решения можно ознакомиться по ссылке.

Теперь все принцессы будут под присмотром :)