![](https://habrastorage.org/getpro/habr/upload_files/59a/d27/630/59ad27630b747ab056a485070ffcd477.png)
Все чаще для сегментации изображений используется глубокое обучение и сверточные нейронные сети. В случае медицинских картинок достаточно сильно проявляются основные проблемы этого метода: не хватает робастности и интерпретируемости. Происходит это в основном из-за того, что CNN обучаются на текстуре изображения, а не на форме, или требуются дополнительные вычисления post hoc, которые, как было показано, ненадежны с точки зрения интерпретируемости.
В статье SAUNet: Shape Attentive U-Net for Interpretable Medical Image Segmentation авторы (Jesse Sun, Fatemeh Darbehani, Mark Zaidi, Bo Wang) предлагают добавить к модели U-Net второй поток данных о форме, а также использовать dual-attention декодер. Такой метод позволил получить очень хорошие результаты на датасетах изображений МРТ сердца SUN09 и AC17, обеспечивая высокую интерпретируемость при различных разрешениях.
Работа опирается на последние достижения в области моделей channel-attention с использованием модулей сжатия и возбуждения, предложенных Hu и др., и spatial attention c оценкой внимания, предложенных Jetley и др..
Суть подхода
ASAUNet состоит из двух потоков: texture stream и gated shape stream.
![](https://habrastorage.org/getpro/habr/upload_files/1e7/237/3dd/1e72373dd9f2576f8d4397c8e2eb3b6a.png)
Texture stream
Текстурный поток имеет структуру U-Net, но энкодер заменен на структуру блоков из DenseNet-121, которые похожи на TiramisuNet. Также авторы предлагают декодер двойного внимания собственной разработки.
Dual attention decoder block
Этот блок объединяет карты, выдаваемые энкодером, с картами декодера с более низким разрешением, которые захватывают больше контекстной и пространственной информации.
Особенно для медицины важно понимать, основываясь на каких особенностях, модель принимает решения.
Авторы предлагают следующую структуру этого блока:
стандартная нормализованная свертка
spatial attention path
channel-wise attention path
Spatial Attention Path
Как раз используется для большей интерпретируемости решений. Состоит из двух сверток , первая из
каналов получает
, вторая 1. Далее применяется сигмоида для вложения в отрезок
. Затем складываем
копий в один канал для дальнейшего умножения.
Channel Attention Path
Состоит из модуля сжатия и внимания, выдает коэффициент сжатия для каждого канала, затем масштабирует на него.
Gated shape stream
Второй поток обрабатывает информацию о форме и границах, используя объекты, обработанные энкодером из первого потока текстур.
На каждом шаге данные из U-Net модели передаются в Shape Stream, где вычисляется карта внимания границ по следующей формуле:
где - сигмоидная функция,
- конкатенация карт по каналам,
и
- feature maps из потоков,
- нормализованная свертка
.
Feature map следующего слоя этого потока получают применением Residual блока, состоящего из двух нормализованных конволюций с skip-connection, к предыдущему, поэлементно домноженному на
:
Функция потерь
Состоит из обычной функции кросс энтропии и dice loss, вычисляющей перекрытие и сходство между двумя наборами.
Первая определена как
Вторая
где- количество классов, а
обозначает
-й пиксель
-го индексированного класса матрицы
.
Также определим - двоичная кросс энтропия предсказанных границ.
Итоговая функция потерь имеет вид
где- гиперпараметры. Авторы используют их равными единице, и все работает хорошо, при этом утверждают, что если занулить один из гиперпараметров, модель плохо и медленно обучается.
Эксперименты
Эксперименты проводились на двух датасетах:
SUN09, два класса (endocardium, epicardium), всего 260 и 135 соответственно двумерных MРT кадров 128 на 128 пикселей. Все модели работали 120 эпох, с размером батча 4.
AC17, 200 наборов МРТ снимков от 100 уникальных пациентов с разрешением 256 на 256, 180 эпох.
Дополнительно авторы провели сравнение на втором датасете, подтверждающее эффективность дополнительного потока.
![](https://habrastorage.org/getpro/habr/upload_files/978/85d/3ab/97885d3ab24dc4255f753f4a8c7e7cfe.png)
Итоговые оценки на датасетах получились следующие:
![](https://habrastorage.org/getpro/habr/upload_files/cd4/160/6dd/cd41606dd435065824012266a5203aff.png)
Интерпретируемось
Для получения карты значимости каждого изображения с использованием модели SmoothGrad приходится прогонять модель 25-50 раз вперед-назад, из-за этого на все 384 изображения ушло 24 минуты. А для модели авторов всего 20 секунд! Кроме этого новый метод предлагает и методы определения значимости на разных уровнях и разрешениях, которые другие модели, основанные на градиенте, не предлагают.
Пример приведен ниже:
![](https://habrastorage.org/getpro/habr/upload_files/daf/a7a/a94/dafa7aa941cdbb898f7dd6f668be4cb0.png)
Заключение
Получилась неплохая модель, особенно учитывая, что она гораздо проще интерпретируемая, чем ее предшественники, поэтому авторы надеются, что приблизили методы глубокого обучения к клиническому применению, обещают проделать дополнительную работу по интерперетируемости.
P.S.
Этот пост написан для https://github.com/spbu-math-cs/ml-course