TransUNet — революция в традиционной сегментации изображений

Обновление U-Net путем объединения CNN и трансформера для достижения результата SOTA в решении задач по сегментации изображений.

Интуиция

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

Однако CNN могут эффективно использовать только информацию ближнего радиуса действия (“последних 100 метров”) (или локальную) из-за небольшого размера ядра свертки, что не позволяет им в достаточной степени исследовать данные в задачах, содержащих компоненты с зависимостями на больших расстояниях.

Трансформеры, обычно используемые для обработки естественного языка (NLP), могут эффективно исследовать информацию на дальних расстояниях, но недостаточно надежны в работе с данными вблизи, как это делают CNN.

Чтобы использовать возможности CNN для компенсации недостатков трансформеров и с другой стороны в задачах сегментации изображений, Чен (Chen) и др. предложили TransUNet, которая также является первой моделью сегментации изображения, построенная на основе трансформера. Стоит также отметить, что авторы удостоверились в перспективности объединения CNN и трансформеров, сначала попытавшись использовать чистую архитектуру трансформера для сегментации изображения. Однако она не сработала так же хорошо, как внедрение CNN в их архитектуру, поскольку трансформеры не так хороши при эксплуатации локальных возможностей, как CNN.

TransUNet

Мы знаем, что в сети U-Net (Рис. 1) есть кодер (даунсемплинг) и декодер (апсемплинг). Даунсемплинг осуществляет кодирование свойств изображения в карту высокого уровня, а ее содержание используется в апсемплинге  для получения конечной маски той же размерности, что и входная.

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

Для разбивки на этапы:

Даунсемплинг (кодирование)

• Во-первых, CNN используется в качестве экстрактора свойств при создании их карты для входного сигнала, как показано в розовой рамке на Рис. 2.

• Для каждого уровня экстрактора свойств их выходная карта (которая кодирует промежуточные карты свойств высокого уровня) затем конкатенируется в маршрут декодера того же уровня, как показано пунктирными стрелками на Рис. 2.

• Затем карта свойств токенизируется (векторизуется) в двумерный эмбеддинг формы (n_patch, D) путем линейного проецирования, причем D — его полная длина. Эмбеддинги были предварительно обучены и будут сохранять позиционную информацию карты свойств (если вы не понимаете, как это делается, не беспокойтесь об этом, так как это не сильно помешает вашему пониманию TransUNet).

• После получения эмбеддингов они подаются на 12 слоев трансформера для кодирования меньшего количества информации на малых расстояниях и большего количества на больших дистанциях от изображения. Каждый слой показан на Рис. 2 (a), где используются модули multi-head self-attention (многоголовое самовнимание) (MSA) и multi-layer perceptron (многослойный персептрон) (MLP). MSA является основным строительным блоком для трансформеров и описан здесь, а MLP просто состоит из нескольких полностью связанных слоев.

• Наконец, чтобы быть готовым к апсемплингу, выходной сигнал преобразуется в (D, H/16, W/16). H/16 и W/16 означают, что высота и ширина уже уменьшились в 16 раз вследствие предыдущих операций.

Апсемплинг (декодирование)

Процесс апсемплинга достаточно прост (без каких-либо причудливых техник).

• Во-первых, входные данные кодера CNN-трансформера обрабатываются сверточным слоем 3×3 с активацией ReLU, апсемплируются, а затем конкатенируются с выходом экстрактора свойств CNN третьего уровня.

• Полученные карты свойств снова проходят через свертку 3×3 со слоем активации ReLU. Затем результаты конкатенируются с данными, полученными с помощью экстрактора свойств CNN второго уровня.

• Шаг повторяется снова. Теперь на выходе получается маска в форме (C, H, W), где C=количество объективных классов, H=высота изображения, а W=ширина изображения.

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

Результаты

Модель была запущена на датасете Synapse Multi-Organ Segmentation (Синапс-полиорганная сегментация). Итоговые результаты оценены по коэффициенту сходства Dice и расстоянию Хаусдорфа.

Как вы можете видеть, TransUNet показала лучшие результаты, чем большинство существующих архитектур SOTA, таких как V-Net, ResNet U-Net, ResNet Attention U-Net и Vision Transformers, что говорит о том, что архитектура на основе Трансформеров лучше использует механизм self-attention, чем другие подобные CNN U-Net.

Реализации

Официальная реализация TransUNet

Уже сегодня вечером состоится открытое занятие «Сверточные нейронные сети». Если интересно, успевайте зарегистрироваться по ссылке.

Также приглашаем всех желающих на вебинар «Компьютерное зрение в спортивной видеоаналитике», на котором:
— Рассмотрим типовые задачи которые ставит бизнес в сфере спортивной видеоаналитики;
— Узнаем об основных подходах и моделях для решения подобных задач;
— Поговорим о трекинге и идентификации игроков, оценке геометрии игровой площадки, а также о распознавании игровых событий и активности игроков на основе видео.
Регистрация доступна по ссылке.