Classification of COVID-19 cases from chest CT volumes using hybrid model of 3D CNN and 3D MLP-Mixer

Resumo Técnico: Classificação de Casos de COVID-19 a partir de Volumes de TC de Tórax Usando um Híbrido de 3D CNN e 3D MLP-Mixer

Definição do Problema
A rápida propagação global da Doença do Coronavírus 2019 (COVID-19) criou uma escassez significativa de mão de obra nas instituições médicas, necessitando de fluxos de trabalho diagnósticos eficientes. Embora a Reação em Cadeia da Polimerase com Transcrição Reversa (RT-PCR) seja a ferramenta diagnóstica padrão, sua sensibilidade é relatada como variando apenas entre 42% e 71%. Em contraste, o diagnóstico baseado em TC de tórax demonstra uma sensibilidade relatada de 97%, tornando-se uma modalidade crítica para a identificação de pneumonia viral. No entanto, a interpretação manual de volumes de TC é demorada. Os sistemas de Diagnóstico Auxiliado por Computador (CAD) automatizados existentes frequentemente dependem de Redes Neurais Convolucionionais (CNNs), que são limitadas por seus campos receptivos e utilizam primordialmente características locais da imagem. Por outro lado, os Vision Transformers (ViT) consideram características globais, mas requerem vastas quantidades de dados de treinamento para alcançar alto desempenho, o que muitas vezes não está disponível em contextos de imagem médica. O desafio reside no desenvolvimento de um método de classificação automatizado que aproveite efetivamente tanto as características locais quanto as globais de volumes de TC 3D, mantendo-se robusto com dados de treinamento limitados.

Metodologia
Os autores propõem um método de classificação automatizada para volumes de TC de tórax utilizando um modelo híbrido que combina uma Rede Neural Convolucional 3D (3D CNN) e uma versão 3D do MLP-Mixer.

• Pré-processamento: Os volumes de TC de tórax de entrada passam por uma segmentação pulmonar para isolar as regiões pulmonares. Fatias axiais que carecem de tecido pulmonar são removidas, e o subvolume restante é escalonado para uma resolução fixa de $192 \times 192 \times 128$ voxels.
• Extração de Características 3D CNN: O subvolume escalonado é processado por uma 3D CNN para extrair características locais da imagem. Esta seção emprega conexões de pooling denso, utilizando uma combinação de pooling máximo e médio (mixed pooling) para reduzir a perda de informação e capturar informações espaciais em múltiplas escalas. O resultado é um volume de características 3D.
• Geração de Patches 3D: Patches 3D não sobrepostos de tamanho $4 \times 4 \times 4$ voxels são extraídos do volume de características 3D. Cada patch é remodelado em um vetor 1D com $C=64$ componentes.
• Classificação 3D MLP-Mixer: Os autores estendem a arquitetura original do 2D MLP-Mixer para lidar com volumes 3D. O conjunto de patches 3D serve como a entrada ($X \in \mathbb{R}^{S \times C}$) para o 3D MLP-Mixer. O modelo processa essas entradas através de $L$ camadas, onde cada camada consiste em dois blocos de Perceptron Multicamadas (MLP):
  1. Mistura de tokens (Token-mixing): Realiza a mistura de características entre patches para integrar o contexto global.
  2. Mistura de canais (Channel-mixing): Realiza a mistura de características intra-patch.
     A saída final é passada por um pooling médio global e camadas totalmente conectadas para produzir uma classificação binária (Alta vs. Baixa probabilidade de COVID-19).

Principais Contribuições

1. Arquitetura Híbrida: O artigo introduz um novo modelo híbrido que integra 3D CNNs para extração de características locais com um 3D MLP-Mixer para integração de características globais. Isso aborda a limitação das CNNs puras (foco local) e dos modelos puramente Transformer/MLP-Mixer (fome de dados).
2. Adaptação do 3D MLP-Mixer: Os autores modificaram com sucesso o MLP-Mixer, originalmente projetado para imagens 2D, para processar dados volumétricos 3D, adaptando os mecanismos de patching e mistura ao espaço 3D.
3. Eficiência com Dados Limitados: A abordagem híbrida é projetada para alcançar alto desempenho de classificação mesmo com um número limitado de amostras de treinamento, uma restrição comum em imagem médica, ao utilizar a 3D CNN para pré-extrair características locais robustas antes da mistura global.

Resultados Experimentais
O método proposto foi avaliado em um conjunto de dados compreendendo 1.205 volumes de TC de tórax (casos de COVID-19 e não-COVID-19) obtidos de múltiplas instituições médicas no Japão. O conjunto de dados foi dividido em 80% para treinamento e 20% para teste.

• Desempenho: O modelo híbrido proposto alcançou uma acurácia de classificação de 79,5%.
• Comparação: Um modelo baseline 3D CNN (consistindo de camadas 3D CNN e camadas MLP simples) treinado no mesmo conjunto de dados alcançou uma acurácia de 74,8%.
• Conclusão: A melhoria de 4,7% na acurácia demonstra que o componente 3D MLP-Mixer contribui efetivamente para o desempenho da classificação ao utilizar características globais da imagem.

Significância e Alegações
O artigo afirma que o uso híbrido de 3D CNN e 3D MLP-Mixer é eficaz para a classificação de volumes 3D, particularmente para pneumonia viral, onde opacidades de vidro fosco e consolidações ocorrem em grandes áreas, exigindo análise de características locais e globais. Os autores enfatizam que esta estrutura híbrida permite alta acurácia sem os massivos conjuntos de dados de treinamento exigidos por modelos puramente MLP-Mixer ou ViT.

No entanto, os autores mantêm um tom modesto em relação à implantação clínica. Eles afirmam que melhorias adicionais na acurácia da classificação são necessárias antes que o modelo possa ser aplicado em ambientes clínicos. O trabalho futuro identificado pelos autores inclui o uso de aumento de dados (data augmentation) para realçar variações anatômicas e de intensidade, a incorporação de conjuntos de dados internacionais e o desenvolvimento contínuo de um sistema completo de CAD para COVID-19. O artigo não afirma que o modelo atual está pronto para uso clínico imediato, mas sim que o apresenta como um passo promissor em direção a esse objetivo.