U-Net: нейросеть для сегментации изображений

U-Net считается одной из стандартных архитектур CNN для задач сегментации изображений, когда нужно не только определить класс изображения целиком, но и сегментировать его области по классу, т. е. создать маску, которая будет разделять изображение на несколько классов. Архитектура состоит из стягивающего пути для захвата контекста и симметричного расширяющегося пути, который позволяет осуществить точную локализацию.

Сеть обучается сквозным способом на небольшом количестве изображений и превосходит предыдущий наилучший метод (сверточную сеть со скользящим окном) на соревновании ISBI по сегментации нейронных структур в электронно-микроскопических стеках. Используя ту же сеть, которая была обучена на изображениях световой микроскопии пропускания (фазовый контраст и DIC), U-Net заняла первое место в конкурсе ISBI 2015 года по трекингу клеток в этих категориях с большим отрывом. Кроме того, эта сеть работает быстро. Сегментация изображения 512×512 занимает менее секунды на современном графическом процессоре.

Для U-Net хатактерно:

• достижение высоких результатов в различных реальных задачах, особенно для биомедицинских приложений;
• использование небольшого количества данных для достижения хороших результатов.

Архитектура U-Net

Архитектура сети приведена на рисунке 1. Она состоит из сужающегося пути (слева) и расширяющегося пути (справа). Сужающийся путь — типичная архитектуре сверточной нейронной сети. Он состоит из повторного применения двух сверток 3×3, за которыми следуют инит ReLU и операция максимального объединения (2×2 степени 2) для понижения разрешения.

На каждом этапе понижающей дискретизации каналы свойств удваиваются. Каждый шаг в расширяющемся пути состоит из операции повышающей дискретизации карты свойств, за которой следуют:

• свертка 2×2, которая уменьшает количество каналов свойств;
• объединение с соответствующим образом обрезанной картой свойств из стягивающегося пути;
• две 3×3 свертки, за которыми следует ReLU.

Обрезка необходима из-за потери граничных пикселей при каждой свертке.

На последнем слое используется свертка 1×1 для сопоставления каждого 64-компонентного вектора свойств с желаемым количеством классов. Всего сеть содержит 23 сверточных слоя.

Обучение U-Net

Сеть обучается методом стохастического градиентного спуска на основе входных изображений и соответствующих им карт сегментации. Из-за сверток выходное изображение меньше входного сигнала на постоянную ширину границы. Применяемая попиксельно, функция soft-max вычисляет энергию по окончательной карте свойств вместе с функцией кросс-энтропии. Кросс-энтропия, вычисляемая в каждой точке, определяется так:

Граница разделения вычисляется с использованием морфологических операций. Затем вычисляется карта весовых коэффициентов:

где wc — карта весов для балансировки частот классов, d1 — расстояние до границы ближайшей ячейки, а d2 — расстояние до границы второй ближайшей ячейки.

Примеры использования и реализации

U-net неоднократно использовалась для сегментации изображений. Ниже приведены некоторые из них:

• классификация нервов с помощью ультразвука на соревновании Kaggle, несколько примеров ядер и образцов кода: 1 2 3 4;
• статья и код по сегментации кровеносных сосудов сетчатки;
• еще одна реализация U-net с Keras;
• применение небольшой сети U-net для распознавания автомобилей.

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

• соревнование по анализу изображений со спутников;
• еще одно соревнование по той же теме;
• финалисты соревнования по раку легких также использовали U-net как часть своей сети.

[Pytorch][Tensorflow][Keras]

Результаты

U-Net применялся в задаче сегментации клеток на изображениях с оптического микроскопа. Эта задача сегментации является частью конкурсов ISBI 2014 и 2015 годов по трекингу клеток.

Датасет PhC-U373 содержит изображения клетки U373 Glioblastoma-astrocytoma на полиакриламидной подложке, записанные методом фазово-контрастной микроскопии. Он содержит 35 частично помеченных изображений для обучения. В этом случае U-Net достигла среднего значения IOU (пересечение по объединению), равного 92%, что значительно лучше предыдущего результата: 83% (см. рис. 2).

Второй датасет DIC-HeLa содержит фотографии клеток HeLa на плоском стекле, записанные с помощью микроскопии дифференциального интерференционного контраста (DIC) (см. рисунки ниже). Он содержит 20 частично помеченных изображений для обучения. В этом случае U-Net достигла среднего значения IOU 77,5%, что значительно лучше, чем последующий алгоритм с его 46%.

Архитектура U-net достигает выдающейся производительности и точности в самых разных приложениях биомедицинской сегментации. Метод требует лишь нескольких помеченных изображений для тренировки и имеет приемлемое время обучения: всего лишь 10 часов на графическом процессоре NVidia Titan (6 ГБ).