fbpx

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

u-net

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

Пример сегментации мозга

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

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

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

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

архитектура u-net
Рисунок 1. Архитектура U-net (пример изображения с разрешением 32×32 пикселя — самым низким). Каждый синий квадрат соответствует многоканальной карте свойств. Количество каналов приведено в верхней части квадрата. Размер x-y приведен в нижнем левом краю квадрата. Белые квадраты представляют собой копии карты свойств. Стрелки обозначают различные операции.

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

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

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

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

Схема U-Net
Схема сети U-net

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

Обучение U-Net

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

кросс-энтропия u-net

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

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

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

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

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

[Pytorch][Tensorflow][Keras]

Результаты

Рисунок 2. Результаты сегментации (IOU) на соревновании ISBI 2015 года по трекингу клеток.

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
Результаты соревнования ISBI по трекингу клеток. (a) Часть входного изображения датасета PhC-U373. (b) Результат сегментации (голубая маска) и метка (желтая граница). (c) Входное изображение датасета DIC-HeLa. (d) Результат сегментации (цветные маски) и метки (желтая рамка).

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