ers_segmentation

Сегментация изображений дистанционного зондирования Земли.

📁 Датасет

Представляет собой папку с изображениями и масками.
Для запуска необходимо скопировать папку с данными в корневую папку проекта.

ML-244-review-task-3
|----- images
        |----- image1.png
                …
        |----- imageN.png
|----- annotations
        |----- mask1.png
                …
        |----- maskN/3.png

Примечательно, что одной маске соответствуют сразу 3 изображения за разные периоды времени.
Посмотрев на данные, соответствующие классы в маске, и сгруппировав их, получаем несколько групп, в которых кол-во представленных классов на маске разное. Наибольшее количество масок содержит сразу все классы. Используем это для разбиения.

Разбиение на train/val/test происходит с помощью:

python prepare_dataset.py –config config.yaml

Данные разбиваются со стратификацией по содержащимся в классам в масках. Также в каждую из выборок попадает маска и все 3 соответствующие ей изображения, чтобы не допустить их появления в другой из выборок.
В результате: 3 csv-файла с полями 'изображение' и 'маска'.

♨️ Аугментации

Для аугментаций использовалась библиотека Albumentations, позволяющая удобно работать с парой изображение-маска и применять соответствующие аугментации к паре.
Аугментации были разделены на 2 категории: для train и для val+test
Тренировочные аугментации содержат: Resize, ShiftScaleRotate, HorizontalFlip, VerticalFlip, OneOf(OpticalDistortion, GridDistortion, GaussianBlur, JpegCompression, RandomBrightnessContrast) с подобранными коэффициентами.

Мозаика

Также для тренировочных данных используется кастомная мозаика: комбинация из патчей 4 рандомно выбранных изображений (объединение 1/4 каждого изображения и соответствующих масок).

⚙️ Конфигурация

Корневая папка проекта содержит файл config.yaml, позволяющий осуществить тонкую настройку для обучения модели. Все параметры названы соответственно.

Модель

В качестве бэкбона используется ResNet34, архитектура UNet.

Функция потерь

Комбинация 2 лоссов с весами:

CrossEntropyLoss (0.7);
DiceLoss (0.3).

Оптимизатор и LR-Scheduler

В качестве оптимизатора: Adam.
LR-Scheduler: CosineAnnealingLR.

Callbacks

ModelCheckpoint для сохранения весов, обеспечивающих наивысший IOU на валидации и EarlyStopping для ранней остановки.

🚀 Тренировка, валидация, тест

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

pip install -r requirements.txt

При наличии GPU необходимо установить torch, torchvision + CUDA.
При использовании CPU (нежелательно): torch, torchvision.

2. Запуск файла:

python train_val_test.py --config config.yaml

✅ Результаты

Метрики аггрегировались с micro-imagewise (с учетом дизбаланса классов для каждого изображения)

Val IOU (micro-imagewise)	Test IOU (micro-imagewise)
0.65	0.67

📝 Выводы

Ключевые моменты в решении данной задачи:

грамотное разбиение на train/val/test;
подбор аугментаций;
мозаика оказала сильное влияние на метрики качества (положительное);
подбор бэкбона и архитектуры;
настройка гиперпараметров тренировки;
комбинация loss-функций оказала положительное влияние на метрики качества.

Что влияет на точность в худшую сторону:

недостаток данных;
простой энкодер;
некоторые аугментации, "портящие" изображения;
неверное разбиение на train/val/test.

Подходы для улучшения качества модели:

использование более сложных моделей/архитектур (например, трансформер в качестве энкодера; DeepLabV3 вместо UNet);
увеличение количества эпох;
увеличение количества данных;
увеличение размера изображений + patching segmentation (скользящее окно и объединение прогнозов);
кастомные аугментации: например, рандомные частичные геометрические преобразования для разных областей-классов;
использование ансамбля моделей мультиклассовой сегментации – отдельная модель для отдельного класса;
использование ансамбля моделей, состоящего из нескольких сетей, при классификации каждого пикселя использовать метод мажоритарного голосования.

licksylick/ers_segmentation