Автоматическая фильтрация изображений животных

Кейс

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

Описание проекта

Нахождение объектов делается с помощью предобученной YOLOv5. Этот проект использует модель Swin Transformer для классификации изображений животных. Для решения проблемы несбалансированных классов была добавлена функция потерь Focal Loss, которая повышает точность классификации при большом дисбалансе между классами. Модель обучается на наборе данных, включающем изображения животных и пустых фонов, с применением техник аугментации и разделения на обучающую и валидационную выборки.

Стек технологий

• PyTorch — для построения и обучения модели
• Hugging Face Transformers — для модели Swin Transformer и предобработки изображений
• Scikit-Learn — для разбиения на обучающую и валидационную выборки
• WandB — для мониторинга и отслеживания метрик
• PIL и OpenCV — для обработки изображений и аугментации

Структура кода

Основные компоненты проекта

1. Загрузка и подготовка изображений
   Для обработки изображений используется AutoFeatureExtractor из библиотеки Hugging Face, который автоматически выполняет базовую предобработку изображений. Дополнительно применяются различные аугментации из torchvision.transforms, чтобы увеличить разнообразие данных и улучшить устойчивость модели.

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

3. Focal Loss
   Реализована специальная функция потерь Focal Loss, которая учитывает дисбаланс между классами, помогая модели лучше обучаться на данных, где некоторые классы встречаются гораздо реже.

4. AnimalTrainer
   Специальный класс AnimalTrainer, созданный на основе Trainer из Hugging Face, в котором переопределен метод compute_loss. Это позволяет использовать Focal Loss вместо стандартной функции потерь, что делает обучение более адаптированным к условиям задачи.

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

6. Аугментация данных
   Для увеличения разнообразия данных применяются следующие аугментации: случайное горизонтальное и вертикальное отражение (RandomHorizontalFlip, RandomVerticalFlip), случайное поворачивание изображений (RandomRotation), случайное изменение размера и обрезка (RandomResizedCrop), а также опциональный эффект размытия движении (MotionBlurTransform). Эти трансформации делают модель более устойчивой к различным условиям съемки.

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

Запуск

Установка зависимостей

pip install -r requirements.txt

Запуск проекта

python main_app.py

Параметры и аргументы обучения

• learning_rate: Скорость обучения установлена на 2e-6 для плавного обучения.
• batch_size: Размер батча для обучения — 20, для валидации — 10.
• num_train_epochs: Обучение проводится в 40 эпох.
• weight_decay: Регуляризация с коэффициентом 0.01 для предотвращения переобучения.

Метрики

Для оценки модели использовался roc_auc.

Примеры работы программы