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

Интуиция

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

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

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

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

TransUNet

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

Рис. 1: Диаграмма U-Net из оригинальной статьи Роннебергера (Ronneberger) и др.
Рис. 1: Диаграмма U-Net из оригинальной статьи Роннебергера (Ronneberger) и др.

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

Рис. 2: Обзор архитектуры TransUNet. MSA (Multi-head Self-Attention. Многоголовое самовнимание), и MLP (Multi-Layer Perceptron. Многослойный персептрон).
Рис. 2: Обзор архитектуры TransUNet. MSA (Multi-head Self-Attention. Многоголовое самовнимание), и MLP (Multi-Layer Perceptron. Многослойный персептрон).

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

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

  • Во-первых, 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 раз вследствие предыдущих операций.

Рис. 3: Упрощенная версия обзора архитектуры TransUNet.
Рис. 3: Упрощенная версия обзора архитектуры TransUNet.

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

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

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

  • Полученные карты свойств снова проходят через свертку 3x3 со слоем активации 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.

Таблица 1: Сводные данные о производительности TransUNet по сравнению с другими моделями сегментации изображений SOTA.
Таблица 1: Сводные данные о производительности TransUNet по сравнению с другими моделями сегментации изображений SOTA.

Реализации

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


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

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

Комментарии (0)