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 ГБ).