Swin-PSAxialNet: uma técnica eficiente de segmentação de múltiplos órgãos

Resumo

O presente protocolo descreve um método eficiente de segmentação de múltiplos órgãos chamado Swin-PSAxialNet, que alcançou excelente precisão em comparação com os métodos de segmentação anteriores. As principais etapas deste procedimento incluem coleta de conjunto de dados, configuração do ambiente, pré-processamento de dados, treinamento e comparação de modelos e experimentos de ablação.

Resumo

A segmentação abdominal de múltiplos órgãos é um dos tópicos mais importantes no campo da análise de imagens médicas e desempenha um papel importante no suporte a fluxos de trabalho clínicos, como diagnóstico de doenças e planejamento de tratamento. Neste estudo, é proposto um método eficiente de segmentação de múltiplos órgãos denominado Swin-PSAxialNet baseado na arquitetura nnU-Net. Ele foi projetado especificamente para a segmentação precisa de 11 órgãos abdominais em imagens de TC. A rede proposta fez as seguintes melhorias em comparação com a nnU-Net. Em primeiro lugar, foram introduzidos módulos Space-to-depth (SPD) e blocos de extração de recursos de atenção axial compartilhada por parâmetros (PSAA), aprimorando a capacidade de extração de recursos de imagem 3D. Em segundo lugar, uma abordagem de fusão de imagens em várias escalas foi empregada para capturar informações detalhadas e características espaciais, melhorando a capacidade de extrair características sutis e características de borda. Por fim, um método de compartilhamento de parâmetros foi introduzido para reduzir o custo computacional e a velocidade de treinamento do modelo. A rede proposta atinge um coeficiente de Dice médio de 0,93342 para a tarefa de segmentação envolvendo 11 órgãos. Os resultados experimentais indicam a notável superioridade do Swin-PSAxialNet sobre os métodos de segmentação convencionais anteriores. O método apresenta excelente precisão e baixos custos computacionais na segmentação dos principais órgãos abdominais.

Introdução

A intervenção clínica contemporânea, incluindo o diagnóstico de doenças, a formulação de planos de tratamento e o rastreamento dos resultados do tratamento, depende da segmentação precisa das imagens médicas¹. No entanto, as complexas relações estruturais entre os órgãos abdominais²tornam uma tarefa desafiadora alcançar a segmentação precisa de vários órgãos abdominais³. Nas últimas décadas, os desenvolvimentos florescentes em imagens médicas e visão computacional apresentaram novas oportunidades e desafios no campo da segmentação abdominal de múltiplos órgãos. A tecnologia de Ressonância Magnética (RM)⁴ e Tomografia Computadorizada (TC)⁵ nos permite adquirir imagens abdominais de alta resolução. A segmentação precisa de múltiplos órgãos a partir de imagens de TC tem valor clínico significativo para a avaliação e tratamento de órgãos vitais, como fígado, rins, baço, pâncreas, etc.6,7,8,9,10 No entanto, a anotação manual dessas estruturas anatômicas, especialmente aquelas que requerem intervenção de radiologistas ou oncologistas de radiação, é demorada e suscetível a influências subjetivas¹¹. Portanto, há uma necessidade urgente de desenvolver métodos automatizados e precisos para segmentação abdominal de múltiplos órgãos.

Pesquisas anteriores sobre segmentação de imagens dependiam predominantemente de redes neurais convolucionais (CNNs), que melhoram a eficiência da segmentação empilhando camadas e introduzindo o ResNet¹². Em 2020, a equipe de pesquisa do Google apresentou o modelo Vision Transformer (VIT)¹³, marcando um exemplo pioneiro de incorporação da arquitetura Transformer ao domínio visual tradicional para uma variedade de tarefas visuais¹⁴. Embora as operações convolucionais possam contemplar apenas informações de recursos locais, o mecanismo de atenção em Transformers permite a consideração abrangente de informações de recursos globais.

Considerando a superioridade das arquiteturas baseadas em Transformers sobre as redes convolucionais tradicionais¹⁵, várias equipes de pesquisa realizaram uma extensa exploração na otimização da sinergia entre os pontos fortes dos Transformers e as redes convolucionais 16,17,18,19. Chen et al. introduziram o TransUNet para tarefas de segmentação de imagens médicas¹⁶, que aproveita os Transformers para extrair recursos globais de imagens. Devido ao alto custo do treinamento de rede e à falha em utilizar o conceito de hierarquia de extração de recursos, as vantagens do Transformer não foram totalmente realizadas.

Para resolver esses problemas, muitos pesquisadores começaram a experimentar a incorporação de Transformers como a espinha dorsal para treinar redes de segmentação. Liu et ^al.17 introduziram o Transformador Swin, que empregava um método de construção hierárquica para extração de características em camadas. O conceito de Windows Multi-Head Self-Attention (W-MSA) foi proposto, reduzindo significativamente o custo computacional, particularmente na presença de mapas de recursos de nível superficial maiores. Embora essa abordagem reduzisse os requisitos computacionais, ela também isolava a transmissão de informações entre diferentes janelas. Para resolver esse problema, os autores introduziram ainda o conceito de Shifted Windows Multi-Head Self-Attention (SW-MSA), permitindo a propagação de informações entre janelas adjacentes. Com base nessa metodologia, Cao et al. formularam o Swin-UNet¹⁸, substituindo as convoluções 2D em U-Net por módulos Swin e incorporando W-MSA e SW-MSA nos processos de codificação e decodificação, alcançando resultados de segmentação louváveis.

Por outro lado, Zhou et al. destacaram que a vantagem da operação conv não pode ser ignorada ao processar imagens de alta resolução¹⁹. Seu nnFormer proposto emprega um método de computação de auto-atenção baseado em blocos de imagem tridimensionais locais, constituindo um modelo Transformer caracterizado por uma estrutura em forma de cruz. A utilização da atenção baseada em blocos tridimensionais locais reduziu significativamente a carga de treinamento na rede.

Dados os problemas com o estudo acima, uma estrutura hierárquica híbrida eficiente para segmentação de imagens médicas 3D, denominada Swin-PSAxialNet, é proposta. Este método incorpora um bloco de redução de amostragem, bloco Space-to-depth (SPD)²⁰ , capaz de extrair informações globais²¹. Além disso, adiciona um módulo de atenção axial compartilhada por parâmetros (PSAA), que reduz a contagem de parâmetros de aprendizado de quadrático para linear e terá um bom efeito na precisão do treinamento da rede e na complexidade dos modelos de treinamento²².

Rede Swin-PSAxialNet
A arquitetura geral da rede adota a estrutura em forma de U do nnU-Net²³, consistindo em estruturas de codificador e decodificador. Essas estruturas se envolvem na extração de características locais e na concatenação de características de imagens de grande e pequena escala, conforme ilustrado na Figura 1.

Figura 1: Diagrama esquemático Swin-PSAxialNet da arquitetura de rede. Clique aqui para ver uma versão maior desta figura.

Na estrutura do codificador, o bloco Conv tradicional é combinado com o bloco SPD²⁰ para formar um volume de downsampling. A primeira camada do codificador incorpora Patch Embedding, um módulo que particiona os dados 3D em patches 3D, (P₁, P₂, P₃) representa patches não sobrepostos neste contexto, significa o comprimento da sequência de patches 3D. Após a camada de incorporação, a próxima etapa envolve uma unidade de redução de amostragem convolucional não sobreposta, compreendendo um bloco convolucional e um bloco SPD. Nessa configuração, o bloco convolucional tem um passo definido como 1 e o bloco SPD é empregado para dimensionamento de imagem, levando a uma redução de quatro vezes na resolução e um aumento de duas vezes nos canais.

Na estrutura do decodificador, cada bloco de upsample após a camada de Feição de Gargalo consiste em uma combinação de um bloco de upsampling e um bloco PSAA. A resolução do mapa de recursos é aumentada duas vezes e a contagem de canais é reduzida pela metade entre cada par de estágios do decodificador. Para restaurar informações espaciais e melhorar a representação de recursos, a fusão de recursos entre imagens de grande e pequena escala é realizada entre os blocos de aumento de amostragem. Por fim, os resultados do upsampling são alimentados na camada Head para restaurar o tamanho da imagem original, com um tamanho de saída de (H × W × D × C, C = 3).

Arquitetura de bloco SPD
Nos métodos tradicionais, a seção de redução da amostragem emprega uma única passada com um tamanho de passo de 2. Isso envolve o agrupamento convolucional em posições locais na imagem, limitando o campo receptivo e confinando o modelo à extração de recursos de pequenos patches de imagem. Este método utiliza o bloco SPD, que divide finamente a imagem original em três dimensões. A imagem 3D original é segmentada uniformemente ao longo dos eixos x, y e z, resultando em quatro corpos de subvolume. (Figura 2) Posteriormente, os quatro volumes são concatenados por meio da operação "cat", e a imagem resultante sofre uma convolução 1 × 1 × 1 para obter a imagem com amostragem^{reduzida 20}.

Figura 2: Diagrama de blocos SPD. Clique aqui para ver uma versão maior desta figura.

Arquitetura de blocos PSAA
Em contraste com as redes tradicionais da CNN, o bloco PSAA proposto é mais eficaz na condução do foco global de informações e mais eficiente no aprendizado e treinamento da rede. Isso permite a captura de imagens e recursos espaciais mais ricos. O bloco PSAA inclui aprendizado de atenção axial com base no compartilhamento de parâmetros em três dimensões: Altura, Largura e Profundidade. Em comparação com o mecanismo de atenção convencional que realiza o aprendizado de atenção para cada pixel da imagem, esse método conduz o aprendizado de atenção de forma independente para cada uma das três dimensões, reduzindo a complexidade da autoatenção de quadrática para linear. Além disso, um mecanismo de compartilhamento de parâmetros de consultas de chaves que podem ser aprendidas é empregado, permitindo que a rede execute operações de mecanismo de atenção em paralelo nas três dimensões, resultando em uma representação de recursos mais rápida, superior e eficaz.

Access restricted. Please log in or start a trial to view this content.

Protocolo

O presente protocolo foi aprovado pelo Comitê de Ética da Universidade de Nantong. Envolve a avaliação inteligente e a pesquisa de dados multimodais não invasivos ou minimamente invasivos adquiridos, incluindo imagens médicas humanas, movimentos de membros e imagens vasculares, utilizando tecnologia de inteligência artificial. A Figura 3 mostra o fluxograma geral da segmentação de múltiplos órgãos. Todos os links necessários são fornecidos na Tabela de Materiais.

Figura 3: Fluxograma geral da segmentação de múltiplos órgãos. Clique aqui para ver uma versão maior desta figura.

1. Recolha de conjuntos de dados

1. Faça o download do conjunto de dados AMOS2022²⁴ (consulte a Tabela de Materiais).
2. Selecione 160 imagens como o conjunto de dados de treinamento e 40 imagens como o conjunto de dados de teste.
3. Execute terminalmente mkdir ~/PaddleSeg/contrib/MedicalSeg/data/raw_data.
4. Abra o caminho de raw_data.
5. Coloque o conjunto de dados no diretório raw_data e garanta a correspondência correta com os nomes dos arquivos nos diretórios "imagesTr" e "labelsTr".
6. Crie dataset.json arquivo.
7. Adicione informações do conjunto de dados, como numTraining, numTest e nome de rótulos , ao arquivo "dataset.json".
   NOTA: O conjunto de dados para este estudo consiste em 200 conjuntos de dados de TC de alta qualidade selecionados do conjunto de dados oficial de 500 CT da AMOS2022 para fins de treinamento e teste. O conjunto de dados abrange 11 órgãos, incluindo baço, rim direito, rim esquerdo, vesícula biliar, esôfago, fígado, estômago, aorta, veia cava inferior, pâncreas e bexiga.

2. Configuração do ambiente

1. Instale o ambiente Paddlepaddle (consulte a Tabela de Materiais) e o CUDA versão²⁵ correspondente.
   NOTA: Recomenda-se utilizar o Paddlepaddle versão 2.5.0 e CUDA versão 11.7.
2. Git clone o repositório PaddleSeg²⁵.
3. Execute terminalmente cd ~/PaddleSeg/contrib/MedicalSeg/.
4. Execute pip install -r requirements.txt.
5. Execute pip install medpy.

3. Pré-processamento de dados

1. Modifique o arquivo "yml", incluindo parâmetros como data_root, batch_size, modelo, train_dataset, val_dataset, otimizador, lr_scheduler e perda.
2. Especifique os parâmetros necessários para a execução, incluindo "-config, --log_iters, --precision, --nnunet, --save_dir, --save_interval, etc.
   NOTA: O "raw_data" original é transformado em três tipos de arquivos de dados: "Cropped", "Decathlon" e "Preprocessed", atendendo a diferentes requisitos de treinamento e avaliação para melhorar o desempenho do modelo e a eficácia do treinamento.

4. Treinamento e comparação de modelos

NOTA: Como uma linha de base amplamente utilizada no campo da segmentação de imagens, o nnU-Net²³ serve como um modelo de linha de base no estudo. O processo de comparação de modelos específicos é o seguinte.

1. Trem nnU-Net
   1. Configure com um tamanho de lote de 2 e 80.000 iterações, registre a cada 20 iterações e utilize a precisão fp16.
   2. Defina o save_dir como a pasta "saída" dentro do diretório de trabalho, com um intervalo de salvamento definido em 1.000 iterações.
   3. Adicione o parâmetro --use_vdl para visualização do parâmetro VisualDL.
   4. Execute train.py.
   5. Crie quatro arquivos yml , especifique as dobras de 1 a 4 e treine o modelo com validação cruzada quíntupla.
2. Trem Swin-UNet
   1. Manter a consistência com as condições experimentais previamente estabelecidas.
   2. Remova o parâmetro --nnunet e use --swinunet.
   3. Execute o arquivo train.py.
   4. Treine dobras de 1 a 4 individualmente e, em seguida, sujeito ao processamento do conjunto.
3. Trem TransUNet
   1. Siga os procedimentos descritos na etapa 4.1.
   2. Abra o arquivo yml .
   3. Modifique type: nnunet para type: transunet.
   4. Modifique o parâmetro --nnunet para --transunet.
   5. Execute o arquivo train.py.
   6. Treine as dobras 1-4 individualmente e processe o processamento do conjunto.
4. Trem UNETR
   1. Mantenha variáveis experimentais consistentes.
   2. Repita os procedimentos descritos anteriormente, como Swin-UNet e TransUNet (etapa 4.2 e etapa 4.3).
   3. Execute train.py.
   4. Treine as dobras 1-4 individualmente e processe o processamento do conjunto.
5. Formador de Trem
   1. Modifique os parâmetros como Swin-UNet e TransUNet.
   2. Execute uma única dobra train.py.
   3. Execute o processamento de conjunto.

5. Experiência de ablação

1. Bloco PSAA para toda a rede
   1. Mantenha as variáveis consistentes com a etapa 4.1.
   2. Conjunto volume_shape = (64.192.192) e Axial_attention = Verdadeiro.
   3. Execute uma única dobra train.py.
   4. Processe o processamento do conjunto em cinco dobras.
2. Bloco SPD para toda a rede
   1. Mantenha as variáveis consistentes com a etapa 4.1.
   2. Defina SPD_enable = Verdadeiro e Axial_attention = Falso.
   3. Execute uma única dobra train.py.
   4. Processe o processamento do conjunto em cinco dobras.
3. Bloco PSAA para aumento da amostragem
   NOTA: Não use o bloco SPD para downsampling.
   1. Use o bloco PSAA para upsampling.
   2. Mantenha as variáveis consistentes com a etapa 4.1.
   3. Defina Axial_attention = False na parte de redução da resolução e SPD_enable = False.
   4. Execute uma única dobra train.py.
   5. Processe o processamento do conjunto de cinco dobras.
4. Bloco PSAA para aumento da amostragem
   NOTA: Use o bloco SPD para downsampling.
   1. Use o bloco SPD para redução da amostragem e o bloco PSAA para upsampling.
   2. Mantenha as variáveis consistentes com a etapa 4.1.
   3. Defina Axial_attention = False na parte de redução da resolução e SPD_enable = True.
   4. Execute uma única dobra train.py.
   5. Processe o processamento do conjunto de cinco dobras.
   6. Defina o número de iterações como 2.00.000 devido aos resultados superiores.

Access restricted. Please log in or start a trial to view this content.

Resultados

Este protocolo emprega duas métricas para avaliar o modelo: Dice Similarity Score (DSC) e 95% Hausdorff Distance (HD95). O DSC mede a sobreposição entre as previsões de segmentação de voxel e a verdade fundamental, enquanto o 95% HD avalia a sobreposição entre os limites de previsão de segmentação de voxel e a verdade fundamental, filtrando 5% dos valores discrepantes. A definição de DSC²⁶ é a seguinte:

Access restricted. Please log in or start a trial to view this content.

Discussão

A segmentação dos órgãos abdominais é um trabalho complicado. Em comparação com outras estruturas internas do corpo humano, como o cérebro ou o coração, a segmentação dos órgãos abdominais parece mais desafiadora devido ao baixo contraste e às grandes mudanças de forma nas imagens de TC^27,28. Swin-PSAxialNet é proposto aqui para resolver este difícil problema.

Na etapa de coleta de da...

Access restricted. Please log in or start a trial to view this content.

Materiais

Name	Company	Catalog Number	Comments
AMOS2022 dataset	None	None	Datasets for network training and testing. The weblink is: https://pan.baidu.com/s/1x2ZW5FiZtVap0er55Wk4VQ?pwd=xhpb
ASUS mainframe	ASUS		https://www.asusparts.eu/en/asus-13020-01910200
CUDA version 11.7	NVIDIA		https://developer.nvidia.com/cuda-11-7-0-download-archive
NVIDIA GeForce RTX 3090	NVIDIA		https://www.nvidia.com/en-in/geforce/graphics-cards/30-series/rtx-3090-3090ti/
Paddlepaddle environment	Baidu	None	Environmental preparation for network training. The weblink is: https://www.paddlepaddle.org.cn/
PaddleSeg	Baidu	None	The baseline we use: https://github.com/PaddlePaddle/PaddleSeg