Tl;dr

Решение задачи сегментации, трекинга и оценки активности поросят командой svinkotrack в рамках хакатона AgroCode Hack.

Для трекигна использовались Scaled-YOLOv4 + SORT. Сегментация проводилась с помощью FPN из библиотеки Segmentation Models Pytorch на кропах боксов, полученных из детекции.

Пробный запуск можно осуществить на Google Colab

Установка и запуск

Предполагается наличие установленной Anaconda.

Клонируем репозиторий:

git clone https://github.com/tvelovraf/Pigs-segmentation-and-tracking.git
cd Pigs-segmentation-and-tracking

Установка пакетов

Создаём и активируем новое окружение:

conda create --name <envname>  
conda activate <envname> 

Устанавливаем Python и необходимые библиотеки:

conda install python=3.7.9
conda install pytorch torchvision cudatoolkit=11.3 -c pytorch

Устанавливаем зависимости из requirements.txt:

pip install -r requirements.txt

Скачиваем отсюда веса для детекции и сегментации и кладём их в корень проекта.

Теперь проект готов к запуску.

Запуск

Файл для запуска проекта — detect.py.
В качестве двух аргументов используются:

• --source <path/to/test_video> путь к тестовому видео;
• --output <path/to/inference/> путь к директории для сохранения результата (по умолчанию создаётся папка ./inference/output/ в корне).

Запускаем проект для тестового видео test_cut.mp4:

python detect.py --source test_cut.mp4 --output ./inference/output

Результат сохраняется в виде видео с наложенными боксами, масками, номерами, активностью животных и их количеством для текущего кадра. Также сохраняется картинка activity.png с графиком активностей свиней за всё время.

Пример результата работы

Описание решения

Тренировочный датасет состоял из десяти неразмеченных минутных видео. Прогнав каждый десятый кадр через предобученную Scaled-YOLOv4, мы получили предарительную разметку для детекции. Затем доразметили вручную и обучили модель. Трекинг осуществлялся с помощью SORT, куда подавались результаты детекции.

Для сегментации использовалась модель FPN из библиотеки Segmentation Models Pytorch. На вход сегментации подавались кропы полученных после детекции боксов. Процесс обучения происходил итеративно путём доразметки наиболее ошибочных масок.