U-NET: КОНВОЛЮЦИОННЫЕ СЕТИ ДЛЯ СЕГМЕНТАЦИИ БИОМЕДИЦИНСКИХ ИЗОБРАЖЕНИЙ

U-NET: CONVOLUTIONARY NETWORKS FOR BIOMEDICAL IMAGE SEGMENTATION

Alekseii Sluzhevskiy

Student, Institute of Cybersecurity and Digital Technologies, Russian Technological University MIREA,

Russia, Moscow

АННОТАЦИЯ

Для успешного обучения глубоких сетей требуется множество аннотированных обучающих образцов. В данной статье представлена сеть и стратегия обучения, использующие активное дополнение данных для эффективного применения ограниченного числа аннотированных изображений. Архитектура включает сужающийся путь для захвата контекста и симметричный расширяющийся путь для точной локализации. Такая сеть может быть обучена на небольшом количестве данных и превосходит предыдущие методы, например, конволюционную сеть со скользящим окном, в задаче сегментации нейронных структур на изображениях электронной микроскопии. Используя ту же сеть для изображений фазового контраста и ДИК, достигнуты победы в конкурсе ISBI 2015 по отслеживанию клеток. Сегментация изображения размером 512x512 выполняется менее чем за секунду на современном графическом процессоре.

ABSTRACT

Successful training of deep networks requires many annotated training samples. This paper presents a network and learning strategy that utilizes active data augmentation to efficiently apply a limited number of annotated images. The architecture includes a narrowing path for context capture and a symmetric expanding path for accurate localization. Such a network can be trained on a small amount of data and outperforms previous methods, such as the convolutional network with sliding window, in the task of segmenting neural structures in electron microscopy images. Using the same network for phase contrast and DIC images, wins in the ISBI 2015 cell tracking competition have been achieved. Segmentation of a 512x512 image is accomplished in less than a second on a state-of-the-art GPU.

Ключевые слова: U-Net, глубокие сети, сегментация.

Keywords: U-Net, deep networks, segmentation.

ВВЕДЕНИЕ

За последние два года глубокие сверточные сети превзошли современные достижения в задачах визуального распознавания [7,3]. Их раннее развитие сдерживали ограничения размеров обучающих наборов и архитектур [8]. Прорыв был достигнут при обучении крупной сети на 1 миллион изображений ImageNet [7], что стимулировало создание ещё более глубоких моделей [12].

Традиционно сверточные сети применяются для классификации изображений, где каждому изображению присваивается одна метка. Однако в задачах, требующих локализации, например в биомедицинской обработке, необходимо определять класс для каждого пикселя. Из-за ограниченного объема доступных данных Сиресан и соавт. [1] предложили метод скользящего окна, используя локальные патчи для предсказания пиксельных меток. Этот подход обеспечил победу в соревновании ISBI 2012 по сегментации ЭМ-изображений, однако обладал недостатками: высокой вычислительной сложностью из-за обработки перекрывающихся патчей и ограниченным балансом между контекстом и точностью локализации.

Современные методы [11,4] предложили использование признаков с нескольких слоев для улучшения локализации без потери контекста. В качестве развития идеи используется полностью сверточная сеть [9], дополненная слоями апсемплинга вместо объединения. Такая архитектура восстанавливает высокое разрешение выхода и сочетает признаки высокого разрешения с увеличенным выходом, позволяя формировать более точные сегментации.

Рисунок 1. Архитектура U-сети (пример для 32x32 пикселей в самом низком разрешении)

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

Рисунок 2. Стратегия «перекрытие-плитка» для бесшовной сегментации произвольных больших изображений

Для решения проблемы нехватки обучающих данных применяется агрессивное увеличение данных посредством эластичных деформаций, что особенно эффективно в биомедицинских приложениях, где ткани подвергаются естественным деформациям. Это позволяет сети учиться инвариантности к подобным изменениям [2].

Для разделения соприкасающихся объектов одного класса используется модифицированная функция потерь с повышенным весом для фона между объектами, что улучшает качество сегментации.

АРХИТЕКТУРА СЕТИ

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

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

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

Для корректной сегментации без швов (см. рисунок 2) размер входных тайлов должен обеспечивать применение всех операций 2×2 max-pooling к слоям с четными размерами по осям x и y.

ОБУЧЕНИЕ

Обучение сети проводится с использованием стохастического градиентного спуска в среде Caffe [6]. Из-за непропадированных сверток выходное изображение получается меньше входного на фиксированную ширину границы. Для эффективного использования памяти GPU предпочтение отдается крупным входным плиткам при минимальном размере пакета, который ограничивается одним изображением. Чтобы стабилизировать обучение при таком подходе, используется высокий коэффициент момента (0,99).

Рисунок 3. Клетки HeLa на стекле, полученные с помощью ДИК-микроскопии (дифференциальный интерференционный контраст)

Функция потерь формируется через послойное применение softmax к выходной карте признаков с последующим расчетом кросс-энтропии. Эта функция штрафует отклонения вероятности правильного класса от единицы на каждом пикселе.

Для учета различной частоты классов и улучшения обучения на разделяющих границах между соприкасающимися объектами вводится специальная карта весов. Карта весов предварительно рассчитывается для каждой разметки. Малые разделяющие границы между объектами формируются с помощью морфологических операций. Весовая карта включает балансировку по частотам классов и дополнительное усиление около границ объектов.

Корректная инициализация весов критически важна для глубоких сетей с множеством путей. Для обеспечения равномерной активации по всем слоям начальные веса выбираются из гауссовского распределения с заданной стандартной девиацией, зависящей от числа входных связей каждого нейрона [5]. Например, для свертки 3×3 и 64 каналов число входов составляет 576.

РАСШИРЕНИЕ ДАННЫХ

Дополнение данных критично для обучения сети инвариантности и устойчивости при небольшом объеме обучающей выборки. В задачах микроскопической съемки особенно важны инвариантность к сдвигам, поворотам, деформациям и изменениям яркости. Ключевым методом является случайная эластическая деформация обучающих образцов, позволяющая эффективно обучать сегментационные сети при малом числе размеченных изображений.

Для создания деформаций применяются случайные векторы смещений на грубой сетке 3×3, значения которых выбираются из гауссовского распределения со стандартным отклонением в 10 пикселей. Деформации в масштабе пикселя рассчитываются с использованием бикубической интерполяции. Дополнительное увеличение разнообразия данных обеспечивается слоями дропаута в конце сужающегося пути.

ЭКСПЕРИМЕНТЫ

Применение U-Net продемонстрировано в трёх различных задачах сегментации. Первая задача — сегментация нейронных структур на электронно-микроскопических изображениях из конкурса EM Segmentation Challenge [14], начатого на ISBI 2012. Набор данных включает 30 изображений (512×512 пикселей) вентрального нервного шнура личинки дрозофилы с полными аннотациями клеток и мембран. Оценка производится на секретном тестовом наборе путём вычисления ошибок искривления, Рэнда и пикселя. U-Net, усреднённая по 7 поворотам данных, без дополнительной обработки достигла ошибки искривления 0,0003529 и ошибки Рэнда 0,0382, превзойдя сеть скользящего окна Сиресан и соавт. [1].

Вторая задача — сегментация клеток на фазово-контрастных изображениях в наборе «PhC-U373», содержащем 35 частично размеченных обучающих изображений. Средний показатель IOU составил 92 %, что значительно выше результата второго лучшего алгоритма (83 %).

Третья задача — сегментация клеток HeLa на стекле с использованием DIC-микроскопии. Набор «DIC-HeLa» содержит 20 частично размеченных изображений; достигнут средний IOU 77,5 %, заметно превосходящий результат второго алгоритма (46 %).

Таблица 1.

Рейтинг на EM segmentation challenge [14], отсортированный по ошибке искривления

Рисунок. 4. Результат решения задачи ISBI по отслеживанию клеток.

 (a) часть входного изображения набора данных «PhC-U373». (b) Результат сегментации (маска голубого цвета) с ручным определением истинности (желтая граница) (c) входное изображение набора данных «DIC-HeLa». (d) Результат сегментации (случайные цветные маски) с ручной проверкой истинности (желтая граница)

Таблица 2.

Результаты сегментации (IOU) на задаче ISBI по отслеживанию клеток

ЗАКЛЮЧЕНИЕ

Архитектура U-Net обеспечивает высокую эффективность сегментации в биомедицинских приложениях. За счёт дополнения данных упругими деформациями требуется минимальное количество аннотированных изображений, а обучение занимает около 10 часов на GPU NVidia Titan (6 ГБ). Реализация на базе Caffe [6] и обученные сети доступны для применения в различных задачах сегментации.

Список литературы:

1. Ciresan, D.C., Gambardella, L.M., Giusti, A., Schmidhuber, J.: Deep neural networks segment neuronal membranes in electron microscopy images. In: NIPS. pp. 2852–2860 (дата обращения: 10.04.2025)
2. Dosovitskiy, A., Springenberg, J.T., Riedmiller, M., Brox, T.: Discriminative unsupervised feature learning with convolutional neural networks. In: NIPS (дата обращения: 10.04.2025)
3. Girshick, R., Donahue, J., Darrell, T., Malik, J.: Rich feature hierarchies for accurate object detection and semantic segmentation. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) (дата обращения: 10.04.2025)
4. Hariharan, B., Arbelez, P., Girshick, R., Malik, J.: Hypercolumns for object segmentation and fine-grained localization (2014), arXiv:1411.5752 [cs.CV]
5. He, K., Zhang, X., Ren, S., Sun, J.: Delving deep into rectifiers: Surpassing humanlevel performance on imagenet classification (2015), arXiv:1502.01852 [cs.CV] (дата обращения: 15.04.2025)
6. Jia, Y., Shelhamer, E., Donahue, J., Karayev, S., Long, J., Girshick, R., Guadarrama, S., Darrell, T.: Caffe: Convolutional architecture for fast feature embedding (2014), arXiv:1408.5093 [cs.CV] (дата обращения: 15.04.2025)
7. Krizhevsky, A., Sutskever, I., Hinton, G.E.: Imagenet classification with deep convolutional neural networks. In: NIPS. pp. 1106–1114 (дата обращения: 15.04.2025)
8. LeCun, Y., Boser, B., Denker, J.S., Henderson, D., Howard, R.E., Hubbard, W., Jackel, L.D.: Backpropagation applied to handwritten zip code recognition. Neural Computation 1(4), 541–551 (дата обращения: 20.04.2025)
9. Long, J., Shelhamer, E., Darrell, T.: Fully convolutional networks for semantic segmentation (2014), arXiv:1411.4038 [cs.CV] (дата обращения: 20.04.2025)
10. Maska, M., (...), de Solorzano, C.O.: A benchmark for comparison of cell tracking algorithms. Bioinformatics 30, 1609–1617 (дата обращения: 24.04.2025)
11. Seyedhosseini, M., Sajjadi, M., Tasdizen, T.: Image segmentation with cascaded hierarchical models and logistic disjunctive normal networks. In: Computer Vision (ICCV), 2013 IEEE International Conference on. pp. 2168–2175 (дата обращения: 24.04.2025)
12. Simonyan, K., Zisserman, A.: Very deep convolutional networks for large-scale image recognition (2014), arXiv:1409.1556 [cs.CV] (дата обращения: 24.04.2025)
13. WWW: Web page of the cell tracking challenge, http://www.codesolorzano.com/ celltrackingchallenge/Cell_Tracking_Challenge/Welcome.html (дата обращения: 24.04.2025)
14. WWW: Web page of the em segmentation challenge, http://brainiac2.mit.edu/ isbi_challenge/ (дата обращения: 24.04.2025)