Swin-PSAxialNet: эффективный метод сегментации нескольких органов

Краткое содержание

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

Аннотация

Абдоминальная полиорганная сегментация является одной из наиболее важных тем в области анализа медицинских изображений и играет важную роль в поддержке клинических рабочих процессов, таких как диагностика заболеваний и планирование лечения. В данном исследовании предложен эффективный метод многоорганной сегментации под названием Swin-PSAxialNet на основе архитектуры nnU-Net. Он был разработан специально для точной сегментации 11 органов брюшной полости на КТ-изображениях. В предлагаемую сеть были внесены следующие улучшения по сравнению с nnU-Net. Во-первых, были представлены модули Space-to-depth (SPD) и блоки извлечения признаков с разделяемым по параметрам осевого внимания (PSAA), что расширило возможности извлечения признаков 3D-изображений. Во-вторых, для получения подробной информации и пространственных особенностей был использован подход слияния многомасштабных изображений, что улучшило возможности извлечения тонких особенностей и краевых особенностей. Наконец, был введен метод совместного использования параметров, чтобы снизить вычислительные затраты модели и скорость обучения. Предложенная сеть достигает среднего коэффициента Dice 0,93342 для задачи сегментации с участием 11 органов. Экспериментальные результаты указывают на заметное превосходство Swin-PSAxialNet над предыдущими методами сегментации. Метод показывает отличную точность и низкие вычислительные затраты при сегментации основных органов брюшной полости.

Введение

Современное клиническое вмешательство, включая диагностику заболеваний, составление планов лечения и отслеживание результатов лечения, опирается на точную сегментацию медицинских изображений¹. Тем не менее, сложные структурные взаимоотношения между органами брюшной полости²делают задачу достижения точной сегментации нескольких органов брюшной полости^{сложной задачей 3}. За последние несколько десятилетий бурное развитие медицинской визуализации и компьютерного зрения открыло как новые возможности, так и вызовы в области сегментации полиорганов брюшной полости. Передовые технологии магнитно-резонансной томографии (МРТ)⁴ и компьютерной томографии (КТ)⁵ позволяют получать изображения брюшной полости с высоким разрешением. Точная сегментация нескольких органов по изображениям КТ имеет важное клиническое значение для оценки и лечения жизненно важных органов, таких как печень, почки, селезенка, поджелудочная железа и т. д.6,7,8,9,10 Однако ручное аннотирование этих анатомических структур, особенно тех, которые требуют вмешательства радиологов или онкологов-радиологов, занимает много времени и подвержено субъективному влиянию ¹¹. В связи с этим существует острая необходимость в разработке автоматизированных и точных методов полиорганной сегментации брюшной полости.

Предыдущие исследования сегментации изображений в основном опирались на сверточные нейронные сети (СНС), которые повышают эффективность сегментации за счет наложения слоев и внедрения ResNet¹². В 2020 году исследовательская группа Google представила модель Vision Transformer (VIT)¹³, которая стала первым примером внедрения архитектуры Transformer в традиционную визуальную область для ряда визуальных задач¹⁴. В то время как сверточные операции могут рассматривать только информацию о локальных признаках, механизм внимания в «Трансформерах» позволяет всесторонне рассматривать информацию о глобальных признаках.

Учитывая превосходство архитектур на основе трансформаторов над традиционными сверточными сетями¹⁵, многочисленные исследовательские группы провели обширные исследования по оптимизации синергии между сильными сторонами трансформаторов и сверточных сетей 16,17,18,19. Чен и др. представили TransUNet для задач^{сегментации медицинских изображений 16}, которые используют трансформеры для извлечения глобальных характеристик из изображений. Из-за высокой стоимости обучения сети и неспособности использовать концепцию иерархии извлечения признаков, преимущества Transformer не были реализованы в полной мере.

Чтобы решить эти проблемы, многие исследователи начали экспериментировать с использованием трансформеров в качестве основы для обучающих сегментационных сетей. Лю и ^др.17 представили трансформатор Свина, который использовал метод иерархического построения для многоуровневого извлечения признаков. Была предложена концепция Windows Multi-Head Self-Attention (W-MSA), значительно снижающая вычислительные затраты, особенно при наличии более крупных карт функций мелкого уровня. Хотя такой подход снизил вычислительные требования, он также изолировал передачу информации между различными окнами. Чтобы решить эту проблему, авторы ввели концепцию Shifted Windows Multi-Head Self-Attention (SW-MSA), позволяющую распространять информацию между соседними окнами. Основываясь на этой методологии, Цао и др. сформулировали Swin-UNet¹⁸, заменив двухмерные свертки в U-Net на модули Swin и включив W-MSA и SW-MSA в процессы кодирования и декодирования, достигнув похвальных результатов сегментации.

И наоборот, Zhou et al. подчеркнули, что преимущество операции conv нельзя игнорировать при обработке изображений с высоким разрешением¹⁹. Предложенный ими nnFormer использует метод вычисления собственного внимания, основанный на локальных трехмерных блоках изображений, составляющих модель Трансформера, характеризующуюся крестообразной структурой. Использование внимания на основе локальных трехмерных блоков значительно снижало учебную нагрузку на сеть.

Учитывая проблемы, связанные с вышеупомянутым исследованием, предлагается эффективная гибридная иерархическая структура для сегментации 3D-медицинских изображений, получившая название Swin-PSAxialNet. Этот метод включает в себя блок понижающей дискретизации, блок Space-to-depth (SPD)²⁰ , способный извлекать глобальную информацию²¹. Кроме того, он добавляет модуль параметра «Разделяемое осевое внимание» (PSAA), который сокращает количество параметров обучения с квадратичного до линейного и окажет хорошее влияние на точность обучения сети и сложность обучающих моделей²².

Сеть Swin-PSAxialNet
Общая архитектура сети использует U-образную структуру nnU-Net²³, состоящую из структур энкодера и декодера. Эти структуры участвуют в локальном извлечении признаков и конкатенации признаков из больших и мелких изображений, как показано на рисунке 1.

Рисунок 1: Принципиальная схема сетевой архитектуры Swin-PSAxialNet. Пожалуйста, нажмите здесь, чтобы просмотреть увеличенную версию этого рисунка.

В структуре энкодера традиционный блок Conv объединен с блоком SPD²⁰ для формирования объема понижающей дискретизации. Первый слой кодировщика включает в себя Patch Embedding, модуль, который разделяет 3D-данные на 3D-патчи , (P₁, P₂, P₃), представляет неперекрывающиеся патчи в этом контексте, обозначает длину последовательности 3D-патчей. Следующий шаг включает в себя неперекрывающуюся сверточную единицу понижающей дискретизации, содержащую как сверточный блок, так и блок SPD. В этой конфигурации сверточный блок имеет шаг, равный 1, а блок SPD используется для масштабирования изображения, что приводит к четырехкратному снижению разрешения и двукратному увеличению каналов.

В структуре декодера каждый блок повышающей дискретизации после уровня Bottleneck Feature состоит из комбинации блока повышающей дискретизации и блока PSAA. Разрешение карты функций увеличивается в два раза, а количество каналов уменьшается вдвое между каждой парой каскадов декодера. Для восстановления пространственной информации и улучшения представления объектов выполняется слияние объектов между большими и мелкомасштабными изображениями между блоками повышающей дискретизации. В конечном счете, результаты повышающей дискретизации подаются в слой Head для восстановления исходного размера изображения с выходным размером (H × W × D × C, C = 3).

Блочная архитектура SPD
В традиционных методах секция понижающей дискретизации выполняется одним шагом с шагом 2. Это включает в себя сверточный пул в локальных положениях изображения, ограничение рецептивного поля и ограничение модели извлечением признаков из небольших участков изображения. В этом методе используется блок SPD, который тонко делит исходное изображение на три измерения. Исходное 3D-изображение равномерно сегментировано по осям x, y и z, в результате чего образуются четыре подобъемных тела. (Рисунок 2) Впоследствии четыре тома объединяются с помощью операции «cat», и полученное изображение подвергается свертке 1 × 1 × 1 для получения изображения²⁰ с пониженной дискретизацией.

Рисунок 2: Блок-схема SPD. Пожалуйста, нажмите здесь, чтобы просмотреть увеличенную версию этого рисунка.

Блочная архитектура PSAA
В отличие от традиционных сетей CNN, предлагаемый блок PSAA более эффективен в проведении глобального информационного фокуса и более эффективен в обучении и обучении сети. Это позволяет получать более насыщенные изображения и пространственные особенности. Блок PSAA включает в себя осевое обучение внимания на основе совместного использования параметров в трех измерениях: высота, ширина и глубина. По сравнению с обычным механизмом внимания, который выполняет обучение внимания для каждого пикселя на изображении, этот метод независимо проводит обучение внимания для каждого из трех измерений, снижая сложность собственного внимания с квадратичной до линейной. Кроме того, используется обучаемый механизм совместного использования параметров ключей-запросов, позволяющий сети выполнять операции механизма внимания параллельно в трех измерениях, что приводит к более быстрому, превосходному и эффективному представлению признаков.

Протокол

Настоящий протокол был одобрен Этическим комитетом Наньтунского университета. Он включает в себя интеллектуальную оценку и исследование полученных неинвазивных или минимально инвазивных мультимодальных данных, включая медицинские изображения человека, движения конечностей и визуализацию сосудов, с использованием технологии искусственного интеллекта. На рисунке 3 показана общая блок-схема многоорганной сегментации. Все необходимые веб-ссылки приведены в Таблице материалов.

Рисунок 3: Общая блок-схема многоорганной сегментации. Пожалуйста, нажмите здесь, чтобы просмотреть увеличенную версию этого рисунка.

1. Сбор наборов данных

1. Загрузите набор данных AMOS2022²⁴ (см. Таблицу материалов).
2. Выберите 160 изображений в качестве обучающего набора данных и 40 изображений в качестве тестового набора данных.
3. Terminally execute mkdir ~/PaddleSeg/contrib/MedicalSeg/data/raw_data.
4. Откройте путь raw_data.
5. Поместите набор данных в директорию raw_data и убедитесь в правильном соответствии с именами файлов в директориях "imagesTr" и "labelsTr".
6. Создайте файл dataset.json .
7. Добавьте сведения о наборе данных, такие как numTraining, numTest и имя меток , в файл dataset.json.
   ПРИМЕЧАНИЕ: Набор данных для этого исследования состоит из 200 высококачественных наборов данных КТ, отобранных из AMOS2022 официального набора данных 500 КТ для целей обучения и тестирования. Набор данных охватывает 11 органов, включая селезенку, правую почку, левую почку, желчный пузырь, пищевод, печень, желудок, аорту, нижнюю полую вену, поджелудочную железу и мочевой пузырь.

2. Настройка окружения

1. Установите среду Paddlepaddle (см. Таблицу материалов) и соответствующую версию CUDA²⁵.
   ПРИМЕЧАНИЕ: Рекомендуется использовать Paddlepaddle версии 2.5.0 и CUDA версии 11.7.
2. Клонирование репозитория PaddleSeg в Git²⁵.
3. Терминально выполните cd ~/PaddleSeg/contrib/MedicalSeg/.
4. Выполните pip install -r requirements.txt.
5. Выполните pip install medpy.

3. Предварительная обработка данных

1. Измените файл "yml", включая такие параметры , как data_root, batch_size, модель, train_dataset, val_dataset, оптимизатор, lr_scheduler и потери.
2. Укажите необходимые параметры для выполнения, в том числе "-config, --log_iters, --precision, --nnunet, --save_dir, --save_interval и т.д.
   ПРИМЕЧАНИЕ: Исходный «raw_data» преобразуется в три типа файлов данных: «Обрезанные», «Десятиборные» и «Предварительно обработанные», удовлетворяющие различные требования к обучению и оценке для повышения производительности модели и эффективности обучения.

4. Обучение и сравнение моделей

ПРИМЕЧАНИЕ: Являясь широко используемым базовым уровнем в области сегментации изображений, nnU-Net²³ служит базовой моделью в исследовании. Процесс сравнения конкретных моделей выглядит следующим образом.

1. Поезд nnU-Net
   1. Конфигурация с размером пакета от 2 до 80 000 итераций, ведение журнала каждые 20 итераций и использование точности fp16.
   2. Установите save_dir в качестве папки «вывода» в рабочем каталоге с интервалом сохранения в 1 000 итераций.
   3. Добавьте параметр --use_vdl для визуализации параметров VisualDL.
   4. Запустите train.py.
   5. Создайте четыре yml-файла , укажите свертки 1-4 и обучите модель с помощью пятикратной кросс-валидации.
2. Поезд Swin-UNet
   1. Поддерживайте согласованность с ранее установленными условиями эксперимента.
   2. Удалите параметр --nnunet и используйте --swinunet.
   3. Запустите файл train.py.
   4. Шлейф складывается по 1-4 по отдельности и затем подвергается ансамблевой обработке.
3. Поезд TransUNet
   1. Следуйте процедурам, описанным в шаге 4.1.
   2. Откройте файл yml .
   3. Измените тип: nnunet на тип: transunet.
   4. Измените параметр --nnunet на --transunet.
   5. Запустите файл train.py.
   6. Тренируйтесь по 1-4 фальца по отдельности и обрабатывайте ансамбль обработки.
4. Поезд UNETR
   1. Поддерживайте согласованность экспериментальных переменных.
   2. Повторите описанные ранее процедуры, такие как Swin-UNet и TransUNet (шаг 4.2 и шаг 4.3).
   3. Запустите train.py.
   4. Тренируйтесь по 1-4 фальца по отдельности и обрабатывайте ансамбль обработки.
5. Поезд nnFormer
   1. Измените параметры, такие как Swin-UNet и TransUNet.
   2. Запустите одинарный фальцевый train.py.
   3. Запустите обработку ансамбля.

5. Эксперимент по абляции

1. Блок PSAA для всей сети
   1. Обеспечьте согласованность переменных с шагом 4.1.
   2. Установите volume_shape = (64,192,192) и Axial_attention = true.
   3. Запустите одинарный фальцевый train.py.
   4. Обрабатывайте ансамбль в пять этапов.
2. Блок SPD для всей сети
   1. Обеспечьте согласованность переменных с шагом 4.1.
   2. Установите SPD_enable = True и Axial_attention = False.
   3. Запустите одинарный фальцевый train.py.
   4. Обрабатывайте ансамбль в пять этапов.
3. Блок PSAA для повышающей дискретизации
   ПРИМЕЧАНИЕ: Не используйте блок SPD для понижающей дискретизации.
   1. Используйте блок PSAA для повышения дискретизации.
   2. Обеспечьте согласованность переменных с шагом 4.1.
   3. Установите Axial_attention = False в части понижения дискретизации и SPD_enable = False.
   4. Запустите одинарный фальцевый train.py.
   5. Обработайте пятикратным ансамблем обработку.
4. Блок PSAA для повышающей дискретизации
   ПРИМЕЧАНИЕ: Используйте блок SPD для понижающей дискретизации.
   1. Используйте блок SPD для понижения дискретизации и блок PSAA для повышения дискретизации.
   2. Обеспечьте согласованность переменных с шагом 4.1.
   3. Установите Axial_attention = False в части даунсемплинга и SPD_enable = True.
   4. Запустите одинарный фальцевый train.py.
   5. Обработайте пятикратным ансамблем обработку.
   6. Установите количество итераций равным 2 00 000 из-за превосходных результатов.

Результаты

Этот протокол использует две метрики для оценки модели: индекс сходства кубиков (DSC) и 95% расстояние Хаусдорфа (HD95). DSC измеряет перекрытие между прогнозами воксельной сегментации и наземной достоверностью, в то время как 95% HD оценивает перекрытие между границами прогноза воксельной сегментации и наземной достоверностью, отфильтровывая 5% выбросов. Определение DSC²⁶ выглядит следующим образом:

(1)

Где T и P представляют основные истинные значения и прогнозируемые значения для всех вокселей соответственно. Определение 95% HD²⁶ выглядит следующим образом:

(2)

Где d ^T,P – 95-й процентиль максимального HD₉₅ между наземным и прогнозируемым вокселями, а d ^T,P – 95-й процентиль максимального HD₉₅ между предсказанными вокселями и наземной истиной.

Swin-PSAxialNet сравнивается с различными основными алгоритмами сегментации нескольких органов в таблице 1. Результаты демонстрируют улучшенную производительность сегментации органов брюшной полости в предложенном алгоритме. Точность Dice метода сегментации превосходит среднее значение Dice других методов на 2-3 процентных пункта и даже превосходит nnFormer на 1 процентный пункт. Swin-PSAxialNet также демонстрирует высокую точность в сегментации границ, со средним коэффициентом HD95 5,63, самым низким среди всех сравниваемых алгоритмов.

Методика		ЛКН	Уровень звукового давления	Лив	Сковорода	Аро	Бла	Sto	Гбд	Поз	Эсо	Ркн	Средняя
													ДСК	ХД95
nnU-Net		96.3	96.8	97.21	85.2	95.22	86.57	91.47	84.7	91.79	82.5	95.44	91.2	7.13
Свин-Унет		96.03	96.68	97.22	85.74	95.3	87.82	91.77	84.64	91.67	83.26	96.14	91.48	9.68
ТрансУНет		96.42	96.74	97.05	85.09	94.86	85.2	91.81	83.33	91.52	82.22	96.12	90.94	13.77
УНЭТР		95.94	96.48	96.85	84.14	94.62	85.01	90.18	81.16	90.79	80.34	95.15	90.06	12.92
nnБывшие		96.59	97.22	97.33	86.79	95.31	89.74	92.82	84.88	92.43	84.24	96.32	92.15	6.59
Свин-PSAxialNet		96.85	97.67	97.64	88.39	95.97	92.03	94.4	87.78	93.15	86.24	96.59	93.34	5.63

Таблица 1: Сравнение основных алгоритмов сегментации. Аббревиатуры обозначают следующие органы: Lkn — левая почка, Spl — селезенка, Liv — печень, Pan — поджелудочная железа, Aro — аорта, Bla — мочевой пузырь, Sto — желудок, Gbd — желчный пузырь, Pos — посткава, Eso — пищевод, Rkn — правая почка, DSC — средний коэффициент Dice, HD95 — среднее расстояние Хаусдорфа 95%.

Более того, внедрение механизма совместного использования параметров и усовершенствование структуры трансформатора обеспечивают значительные преимущества как в усложнении модели, так и в скорости вывода (табл. 2, рис. 4 и рис. 5).

Модели	Параметры (M)	ФЛОПС (г)	Время вывода (с)
nn-Unet	17.96	398.54	9.07
УНЭТР	89.58	39.67	12.51
nnБывшие	134.78	152.43	11.24
Свин-PSAxialNet	37.64	43.15	10.31

Таблица 2: Сравнительная таблица показателей модели сети.

Рисунок 4: Кривая вариации кубиков обучающей сети модели. (A) Текущая сеть. (B) nnБывший. (C) nnU-Net. D) ЮНЭТР. Пожалуйста, нажмите здесь, чтобы просмотреть увеличенную версию этой цифры.

Рисунок 5: Сравнение результатов обучения модели. На рисунке показаны результаты текущей сети, результаты nnFormer, результаты nnU-Net и результаты UNETR. Пожалуйста, нажмите здесь, чтобы просмотреть увеличенную версию этой цифры.

В этом исследовании были проведены комплексные эксперименты по абляции на сети алгоритмов, используя блок PSAA и блок SPD в качестве переменных, а также сравнивая однократное обучение с пятикратным обучением. Коэффициенты Дайса из экспериментальных результатов представлены в таблице 3. Эти результаты демонстрируют, что использование блока SPD для понижающей дискретизации и блока PSAA для повышающей дискретизации эффективно повышает производительность сегментации по сравнению с немодифицированной nnU-Net (рис. 6).

Модели	Блок PSAA для всей сети	Блок SPD для всей сети	Блок PSAA только для повышающей дискретизации	Блок PSAA для повышения дискретизации, используйте SPD
Однократное обучение	92.17	90.51	92.24	92.39
Пятикратное обучение	92.82	91.26	92.79	93.34

Таблица 3: Результаты эксперимента по абляции Swin-PSAxialNet. Представлены результаты эксперимента по абляции Swin PSAxialNet. В эксперименте по абляции сравнивали использование блока PSAA и блока SPD, а результаты обучения DSC приведены в таблице.

Рисунок 6: Результаты сегментации. Пожалуйста, нажмите здесь, чтобы просмотреть увеличенную версию этого рисунка.

Обсуждение

Сегментация органов брюшной полости – сложная работа. По сравнению с другими внутренними структурами человеческого тела, такими как мозг или сердце, сегментация органов брюшной полости кажется более сложной задачей из-за низкого контраста и больших изменений формы на КТ-изображениях^27,28. Swin-PSAxialNet предлагается здесь для решения этой непростой проблемы.

На этапе сбора данных это исследование загрузило 200 изображений с официального сайта AMOS2022²⁴. На этапе настройки среды экспериментальная настройка включала среду Paddlepaddle вместе с соответствующей версией CUDA. На этапе предварительной обработки данных была применена случайная обрезка и случайное переключение для повышения надежности данных. На этапах обучения и сравнения модели, основанных на базовой модели nnU-Net, в эксперименте был добавлен блок SPD для даунсемплинга и блок извлечения признаков PSAA для апсемплинга, что позволило получить отличные результаты сегментации. Swin-PSAxialNet сравнивался с различными основными моделями сегментации и продемонстрировал многообещающие результаты на общедоступном наборе данных AMOS2022. Он достиг более высокой точности сегментации по сравнению с другими методами сегментации, упомянутыми в этом исследовании 16,17,18,19. В разделе эксперимента по абляции использование блока PSAA и блока SPD варьировалось для поиска оптимальной сегментационной сети. Результаты эксперимента по абляции показали, что использование модуля PSAA только для повышения дискретизации достигает на один процентный пункт больше, чем использование его для всей сети.

Инновации этого исследования заключаются в следующем: во-первых, модуль SPD встраивается в сеть и образует блок понижающей дискретизации наряду с блоком Conv, тем самым расширяя возможности сети по извлечению признаков. Во-вторых, внедрение блока извлечения признаков PSAA эффективно снижает сложность сети при одновременном повышении точности обучения. Наконец, внедрение механизма совместного использования параметров значительно снижает сложность обучающей модели.

Это исследование имеет ограничения в следующих аспектах: Оно было применено только к AMOS2022 набору данных, и результаты не были проверены на других наборах данных. Опять же, несмотря на то, что nnU-Net достиг высокой производительности в задачах сегментации органов брюшной полости после серии оптимизаций в качестве автоматизированной среды сегментации, некоторые параметры все еще имеют определенную степень субъективности, такие как начальная скорость обучения, количество итераций, архитектура сети и т. д. В будущих исследованиях будут изучены методы автоматического вычисления этих параметров²⁹.

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

Материалы

Список материалов, использованных в этой статье

Имя	Компания	Каталожный номер	Комментарии
AMOS2022 набор данных	Никакой	Никакой	Наборы данных для обучения и тестирования сети. Веб-ссылка: https://pan.baidu.com/s/1x2ZW5FiZtVap0er55Wk4VQ?pwd=xhpb
Основной блок ASUS	ASUS		https://www.asusparts.eu/en/asus-13020-01910200
CUDA версии 11.7	НДВО		https://developer.nvidia.com/cuda-11-7-0-download-archive
Видеокарта NVIDIA GeForce RTX 3090	НДВО		https://www.nvidia.com/en-in/geforce/graphics-cards/30-series/rtx-3090-3090ti/
Окружающая среда для паддлов	Байду	Никакой	Экологическая подготовка к сетевому обучению. Веб-ссылка: https://www.paddlepaddle.org.cn/
ВеслоСег	Байду	Никакой	Базовый уровень, который мы используем: https://github.com/PaddlePaddle/PaddleSeg