U-net для решения семантической сегментации
Предисловие
Ранее в обзорах статей мы рассматривали архитектуры нейронных сетей для решения задачи классификации. Сейчас мы предлагаем вам разобраться с архитектурой U-net для решения задачи семантической сегментации. Эта информация будет полезной, поскольку такая задача (в различных своих вариантах) присутствует в арсенале любого CV инженера.
В основе нашего поста — оригинал статьи “U-Net: Convolutional Networks for Biomedical Image Segmentation”. В современных библиотеках реализация архитектуры может отличаться (причем существенно).
Задача сегментации
Давайте сначала разберемся с понятиями:
- Задача семантической сегментации — задача, при которой необходимо определить класс для каждого пикселя. Получается, на выходе из нейронной сети нужно получить распределение классов для каждого пикселя.
- Задача instance segmentation — задача, при которой необходимо не просто классифицировать принадлежность пикселя к определенному классу, но и определить, к какому именно объекту принадлежат эти пиксели.
Вторая задача сложнее первой: она требует не только сегментации объектов на изображении, но и различения экземпляров одного и того же класса объектов. Поэтому для задачи instance segmentation обычно используются более сложные модели (например, Mask R-CNN).
На рисунке ниже представлен пример маски (segmentation map) для задачи семантической сегментации, где каждый пиксель отнесен к определенному классу. Данные взяты с https://www.cityscapes-dataset.com.
Ключевая идея
Авторы статьи реализовали архитектуру из двух частей: contracting path (левая часть) и expansive path (правая часть). Давайте рассмотрим на изображении ниже схематическое представление архитектуры U-net:
На этом рисунке мы видим: в левой части схемы располагается обычная сверточная нейронная сеть — encoder, а в правой части — decoder. И каждый промежуточный результат encoder’a перед операцией max pooling конкатенируется (skip-connection) с промежуточным результатом правой части архитектуры. Это решает проблему потери информации из-за операции max-pooling.
Архитектура
В оригинальной статье нейронная сеть обучалась на одноканальных биологических снимках, поэтому входное изображение было размером 572х572х1.
При обучении на цветных данных (RGB) входной тензор будет иметь количество каналов, равное 3-ем.
Рассмотрим рисунок выше. На нем мы видим получение тензора 570х570х64 после входа изображения 572х572х1 и операции свертки 3х3 с ReLU. Авторы статьи получают тензоры такого размера, потому что не используют unpadded convolutions (или, по-другому, zero padding) для сохранения исходного размера тензора. Это сохранение помогает не потерять информацию на границах. Далее применяется max pooling, который уменьшает размер тензора вдвое по ширине и высоте. И таким образом мы продолжаем применять conv 3х3 и max polling 2х2, пока тензор у нас не станет размером 28х28х1024.
При применении max pooling теряется важная высокоуровневая информация. Следовательно, точного результата сегментации объектов из тензора 28х28х1024 не выйдет, так как upsampling не сможет ее восстановить. Раньше применяли разные подходы к решению задачи сегментации, например, интерполировали тензор 28х28х1024 к целевому размеру маски. Это давало неточный результат из-за потери контуров объектов (как на рисунке ниже).
Теперь необходимо привести тензор, полученный с помощью encoder, к целевому размеру. Для этого используется up-conv 2х2.
В качестве up-conv могут быть применены различные операции, например, такие как TransposedConv или DilatedConv. Если вы хотите вспомнить об этом, то рекомендуем ознакомиться со статьей про 2D свертки 🙂 В результате такой операции тензор в два раза увеличится по высоте и ширине, а количество фильтров в два раза уменьшится.
Однако сейчас TransposedConv применяется редко, а вместо нее часто используют билинейную интерполяцию и следующую за ней свертку 1х1 с необходимым количеством каналов на выходе.
Затем авторы статьи берут тензор 64х64 с предыдущего этапа и конкатенируют его с результатом операции up-conv (этот процесс показан серой стрелочкой на рисунке ниже). Однако из-за несовпадения размерностей они берут от тензора кроп размером 56х56 (на изображении он обведен пунктирной линией). Такая операция помогает восстановить информацию о фичах, потерянных на этапе их сжатия до тензора размером 28х28. Например, на верхних слоях она позволяет сегментировать края объектов, так как в начальных тензорах encoder’a намного больше информации о мелких деталях объекта.
Можно встретить и другие реализации, где вместо кропа используется resize или padding до нужного размера. На самом деле достаточно подать картинку, размер которой делится на 2^k (где k — количество даунсемплингов в два раза). В нашем случае можно взять количество, равное 4-ем, при условии отсутствия сверток со страйдами больше 1 в encoder. Тогда вышеописанная стратегия с кропом окажется ненужной.
Про рисунок 6 важно отметить следующее: при возвращении к обратному размеру выходной тензор имеет размер не 572х572, а 388х388. Дело в том, что при применении свертки края изображения плохо сегментируются, поскольку по краям свертка проходится только один раз. Для точности сегментирования краев на изображении авторы статьи предложили такой подход, как Overlap-tile strategy. Он заключается в добавлении padding’a в виде зеркальных краев изображения. Отметим также, что данная стратегия обычно не используется — достаточно применить unpadded convolutions. В современных имплементациях входной размер изображения равен выходному.
Исходное изображение в статье имело размерность 512х512, то есть входное изображение расширилось на 30px с каждой стороны. Эта стратегия помогает сохранить информацию по краям изображения. Но при таком подходе, плюс без использования unpadded convolutions, авторы статьи получали изображения из исходной 512х512 в 388х388.
Часто для решения задачи сегментации данные имеют высокое разрешение. Если подавать на вход сети изображения большого размера — они не поместятся в память GPU. Поэтому одним из вариантов решения такой проблемы является разбиение изображения на патчи и сегментация этих патчей частями. Авторы статьи предложили обработку патчами для изображений с высоким разрешением.
На рисунке 7 слева показано изображение с примененной стратегией зеркального отображения краев и обработка этого изображения патчами. Желтая отметка — край исходного изображения, к которому нужно было получить маску. Синяя отметка — изображение, подающееся на вход нейронной сети. После обработки всех частей мы можем собрать полное изображение с наложением. Это дополнительно позволит нам уточнить края для каждого патча.
Для решения задачи сегментации нужно намного меньше данных. Например, авторы статьи проводили обучение всего на 30 изображениях, но с применением аугментации.
В архитектуре U-net не используются полносвязные слои. Это значит, что после обучения на небольших изображениях мы сможем подать на инференс изображение большего размера. Единственное ограничение для нас — память нашей GPU.
После нескольких сверток, в том числе и после свертки 1х1 (это самая крайняя зеленая стрелочка вправо на рисунке 6), которая используется для сопоставления каждого 64 компонентного вектора признаков с желаемым количеством классов, авторы статьи получают выходной тензор 388х388х2. Количество каналов тензора равно двум, так как классов в решаемой задаче было два.
В настоящее время в качестве encoder’a используют предобученные бекбоны, такие как ResNet, EfficientNet, SE-Net и другие, поскольку они извлекают из данных более качественные признаки, а это значительно повышает качество сети в поставленной задаче. Что касается decoder-части: от реализации к реализации существуют отличия. Например, симметричность может сохраняться у U-net с бекбоном ResNet, и при этом не сохраняться у U-net с бекбоном efficientnet.
Также результаты сегментации можно улучшить постпроцессингом CRF или DenseCRF.
Ссылки на дополнительные источники
Для более подробного погружения в тему предлагаем к прочтению следующие источники:
