Семантическая сегментация

Задача

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

Попробуем создать и обучить две модели для семантической сегментации:

• SegNet, часто используемую на практике модель, базирующуюся на архитектуре VGG16 для формирования энкодера;
• U-Net, модель, изначально создававшуюся и оптимизированную для семантической сегментации медицинских изображений.

В качестве эксперимента реализуем также несколько модификаций оригинальной модели U-Net: заменим некоторые слои на их аналоги и посмотрим, скажется ли это на результате.

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

Ноутбук

(ноутбук проекта)

Выводы

В данном проекте мы решали задачу сегментации медицинских снимков. В нашем распоряжении датасет ADDI project. В нём содержатся фотографии различных поражений кожи: меланомы и родинок.

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

Каждая модель обучалась в течение 40 эпох с оптимизатором AdamW. В качестве метрики мы использовали коэффициент Жаккара (IoU).

Сведём данные по значениям метрики IoU на последних эпохах обучения для всех функций потерь.

Видно, что для всех функций потерь удалось достичь достаточно неплохих значений метрики IoU. Наилучшие результаты показала модель с функцией потерь WSR loss, однако результаты в целом различаются не очень сильно.

Из построенных в рамках проекта графиков видно, что 40 эпох хватило, чтобы достичь сходимости для всех функций потерь. При этом стоит отметить, что для функции потерь Dice loss значения на графиках сходились стабильнее всего.

Также мы обучили несколько моделей с различными вариациями архитектуры U-Net.

Первую модель мы реализовали в соответствии с оригинальной статьёй. Для пулинга (уменьшения размера изображений внутри сети) использовался MaxPooling, для апсамплинга (увеличения размера изображений внутри сети) - слои с обратной свёрткой (ConvTranspose).

Во второй модели для пулинга также использовался MaxPooling, а для апсмаплинга был применён nearest-neighbor Upsampling.

В третьей модели для пулинга были применены свёрточные слои.

Все три модели обучались с функцией потерь BCE. Для каждой модели рассчитывался коэффициент IoU.

Сведём результаты экспериментов с U-Net в единую таблицу.

Видно, что в рамках наших экспериментов значения метрики IoU для различных модификаций архитектуры U-Net несколько выше, чем для модели SegNet. Наилучшего значения метрики на валидации удалось добиться при использовании свёрточного пуллинга.

Модели на основе U-Net показывают достаточно высокие значения метрики IoU, но они сходятся несколько хуже моделей на основе архитектуры SegNet. Наилучшую сходимость показала модель со свёрточным пуллингом.

Из проведённых экспериментов видно, что на результаты модели оказывают существенное влияние различные факторы, в том числе выбор функции потерь и нюансов архитектуры. При их выборе следует учитывать несколько факторов: наилучшие значения метрики, стабильность сходимости модели, скорость обучения (количество обучаемых параметров модели). В рамках нашей задачи, если предположить, что заказчику важнее всего значение метрики IoU, наилучший результат продемонстрировала модель на основе архитектуры U-Net со свёрточным пуллингом.

В качестве модели мы выбрали ResNext. Данная модификация архитектуры RexNet была предложена в 2017 году и продемонстрировала хорошие результаты на соревновании ImageNet. В качестве направления дальнейшего развития данного проекта можно попробовать обучить модели с различными архитектурами (DenseNet, ElasticNet и т.д.) и сравнить их результаты.

Обучение происходило с оптимизатором AdamW в течение 20 эпох. В качестве функции потерь была выбрана кросс-энтропия.

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

После обучения мы получили предсказания на тестовых данных и загрузили их на платформу kaggle. Результирующее значение метрики F1 на тестовой выборке (0.993) соответствует максимальному балу за данную задачу.