Improving Prostate Gland Segmentation Using Transformer based Architectures

Título: Melhoria da Segmentação da Próstata Utilizando Arquiteturas Baseadas em Transformers

1. Problema e Motivação

A segmentação automática da glândula prostática em imagens de Ressonância Magnética (MRI) ponderadas em T2 enfrenta dois desafios principais:

• Variabilidade Inter-observador: Diferentes radiologistas (leitores) delineiam os contornos da próstata de maneiras distintas, criando ruído nos dados de treinamento e dificultando a generalização dos modelos.
• Deslocamento de Domínio e Desequilíbrio de Classes: As glândulas prostáticas ocupam menos de 3% do volume da imagem, criando um desequilíbrio severo entre o primeiro plano (próstata) e o fundo. Além disso, modelos baseados em Redes Neurais Convolucionais (CNNs), como o 3D U-Net, têm dificuldade em capturar dependências de longo alcance e contexto global, essenciais para lidar com a heterogeneidade de formas e tamanhos da próstata.

O objetivo do estudo foi investigar se modelos baseados em Transformers (especificamente UNETR e SwinUNETR) podem superar essas limitações, mantendo a precisão e a robustez frente à variabilidade entre leitores e ao tamanho da glândula, comparados a um modelo baseline de 3D U-Net.

2. Metodologia

• Dados:

  • Utilizou-se um conjunto de dados de 546 volumes de MRI T2-weighted da próstata.
  • Os dados foram anotados por dois radiologistas independentes (Leitor #1: 342 volumes; Leitor #2: 204 volumes), com 204 volumes sobrepostos para comparação direta.
  • As anotações foram feitas no desafio público ProstateX (TCIA).

• Arquiteturas Comparadas:

  1. 3D U-Net (Baseline): Arquitetura CNN tradicional com conectividade local.
  2. UNETR: Substitui o encoder da U-Net por um Vision Transformer (ViT) puro, utilizando atenção global. Sofre de crescimento quadrático de memória.
  3. SwinUNETR: Utiliza um encoder baseado no Swin Transformer com atenção em janelas deslocadas (shifted-window). Isso permite capturar contexto global com crescimento de memória quase linear, sendo mais eficiente computacionalmente.

• Estratégias de Treinamento e Avaliação:
  Foram testadas três estratégias para avaliar robustez e generalização:

  1. Treinamento em Coorte Única: Treinar com dados de um leitor e testar no outro (para avaliar sensibilidade ao estilo de anotação).
  2. Treinamento Misto com Validação Cruzada (5-fold): Misturar dados de ambos os leitores para avaliar generalização cruzada.
  3. Segmentação Baseada no Tamanho da Glândula: Estratificar os dados em "grandes" e "pequenas" glândulas e testar a robustez do modelo sob desequilíbrio extremo de classes (subamostragem de dados).

• Otimização e Métricas:

  • Hiperparâmetros: Otimizados automaticamente usando o framework Optuna (taxa de aprendizado, tamanho do lote, tamanho do recurso).
  • Função de Perda: Combinação de Dice Loss ponderada e Binary Cross-Entropy para lidar com o desequilíbrio de classes.
  • Métrica Principal: Coeficiente de Semelhança de Dice (DSC).
  • Análise Estatística: Teste de Wilcoxon e teste de Friedman para verificar significância estatística das diferenças de desempenho.

3. Contribuições Principais

• Primeira Análise Sistemática: Este trabalho realiza o primeiro benchmark detalhado e comparativo de UNETR e SwinUNETR em um único conjunto de dados multi-leitor para segmentação da próstata.
• Mitigação de Ruído de Rótulo: Demonstra que os mecanismos de atenção global e de janelas deslocadas dos Transformers mitigam o ruído causado pela variabilidade inter-observador melhor do que as CNNs.
• Robustez ao Desequilíbrio: Evidencia a superioridade dos Transformers (especialmente SwinUNETR) em cenários de desequilíbrio de classes extremos (glândulas pequenas ou dados limitados de um leitor específico).
• Otimização Automatizada: Aplicação rigorosa de busca de hiperparâmetros via Optuna para cada arquitetura e conjunto de dados, garantindo que as comparações sejam justas e otimizadas.

4. Resultados

• Desempenho Geral (Single Reader Training):

  • O SwinUNETR superou consistentemente os outros modelos.
  • No treinamento com Leitor #1 e teste em Leitor #2, o SwinUNETR atingiu um DSC médio de 0.860, enquanto o UNETR atingiu 0.833 e o U-Net 0.851.
  • O UNETR mostrou-se menos estável em cenários de teste cruzado em comparação ao SwinUNETR.

• Treinamento Misto (Cross-Validation):

  • O SwinUNETR alcançou os melhores resultados: 0.8583 (Leitor #1) e 0.8678 (Leitor #2).
  • O UNETR também melhorou significativamente com dados mistos (0.8517 e 0.8640), tornando-se comparável ao SwinUNETR em alguns cenários, mas o Swin manteve maior consistência.
  • O U-Net apresentou desempenho inferior, especialmente no Leitor #1 (0.826).

• Robustez por Tamanho da Glândula:

  • Glândulas Grandes: O SwinUNETR atingiu um pico de 0.902 de DSC, superando o U-Net (0.846) e o UNETR.
  • Glândulas Pequenas: O SwinUNETR manteve alta estabilidade (DSC ~0.85) mesmo com redução drástica de dados de treinamento (até 10% de um leitor). O U-Net e o UNETR sofreram quedas mais acentuadas de desempenho sob desequilíbrio extremo.

• Análise Estatística:

  • O teste de Friedman confirmou diferenças estatisticamente significativas (p < 0.0001) entre os modelos no Leitor #1, favorecendo o SwinUNETR.
  • Visualmente, o SwinUNETR produziu contornos mais suaves e contínuos, com menos falsos positivos e artefatos de super/sub-segmentação em comparação ao UNETR e U-Net.

5. Significado e Conclusão

O estudo conclui que as arquiteturas baseadas em Transformers, particularmente o SwinUNETR, superam os modelos baseados em CNNs (3D U-Net) para a segmentação da próstata em MRI.

• Vantagens Clínicas: O SwinUNETR oferece um aumento de até 5 pontos percentuais no score de Dice em comparação ao baseline, além de ser robusto à variabilidade de anotação entre diferentes radiologistas e instituições.
• Viabilidade Clínica: Devido ao seu equilíbrio entre precisão e eficiência computacional (graças ao mecanismo de janelas deslocadas), o SwinUNETR é capaz de gerar contornos automatizados em GPUs padrão em tempo quase real, tornando-o uma escolha viável para implementação clínica rotineira.
• Generalização: A capacidade de lidar com desequilíbrio de classes e ruído de rótulo sugere que modelos com atenção global podem ser a base para sistemas de segmentação "agnósticos ao leitor", reduzindo a necessidade de retreinamento extensivo para cada novo centro hospitalar.

Em suma, o trabalho valida a transição de CNNs puras para arquiteturas híbridas e baseadas em Transformers como o novo estado da arte para segmentação de órgãos complexos e de pequeno volume em imagens médicas 3D.