NimbleReg: A light-weight deep-learning framework for diffeomorphic image registration

1. Formulação do Problema

O registro de imagens não linear é crítico para a imagem médica, particularmente na neuroanatomia computacional para alinhar varreduras cerebrais. Embora a Aprendizagem Profunda (DL) tenha revolucionado este campo ao permitir inferência quase instantânea, os métodos existentes baseados em DL sofrem de limitações significativas:

• Alto Custo Computacional: A maioria dos métodos depende de representações de grade 3D densas para imagens de entrada, mapas de características e campos de deformação, levando a um consumo massivo de memória (frequentemente exigindo >24GB de RAM) e requisitos de hardware pesados.
• Falta de Garantias Difeomórficas em Métodos de Superfície: Embora existam métodos leves que utilizam superfícies de contorno de regiões segmentadas para reduzir a complexidade, eles tipicamente carecem de um quadro teórico para fundir múltiplos mapeamentos regionais em uma única transformação difeomórfica globalmente consistente (um mapeamento suave, invertível que preserva a topologia).

Os autores visam desenvolver um quadro de registro DL leve, rápido e difeomórfico que aproveite representações de superfície de estruturas anatômicas em vez de grades volumétricas densas.

2. Metodologia: NimbleReg

O NimbleReg reformula o registro de imagens como o alinhamento de superfícies de contorno de múltiplas regiões anatômicas segmentadas. O quadro consiste em três etapas principais:

A. Pré-processamento

1. Extração de Superfície: Usando ferramentas como Marching Cubes, superfícies são extraídas de regiões segmentadas. Crucialmente, interfaces compartilhadas entre regiões adjacentes são processadas para garantir coordenadas de pontos idênticas em limites compartilhados.
2. Processamento de Malha: Superfícies individuais são mescladas em uma única malha, suavizadas para preservar interfaces e, em seguida, re-divididas em superfícies anatômicas individuais.
3. Pré-alinhamento: Uma transformação poliafina (usando Polaffini) alinha a imagem móvel a uma referência padrão baseada em centróides de região. Isso reduz o espaço de busca para o modelo DL, eliminando a necessidade de aprender pose global ou deslocamentos de baixa frequência.

B. Estimativa de Velocidade (O Modelo de Aprendizagem Profunda)

• Arquitetura: O modelo central, $f_\theta$, utiliza uma base PointNet. Processa nuvens de pontos (superfícies móvel e de referência) usando uma Rede Perceptron Multicamada (MLP) compartilhada para extrair características locais, que são agregadas via max-pooling para criar características globais invariantes a permutações.
• Saída: Em vez de deformar pontos diretamente, o modelo produz vetores de velocidade ($v$) associados aos pontos móveis. Isso trata a saída como uma parametrização esparsa de um campo de velocidade em vez de um campo de deslocamento direto.
• Treinamento: O modelo é treinado de maneira auto-supervisionada.
  • Função de Perda de Ajuste ($L_{fit}$): Usa a distância de Chamfer simétrica para minimizar a lacuna entre pontos móveis transformados e pontos de referência.
  • Função de Perda de Regularização ($L_{reg}$): Penaliza mudanças abruptas de velocidade entre pontos dentro do mesmo simplex (face triangular) da malha de superfície para garantir suavidade ao longo da superfície.

C. Integração em Transformação Difeomórfica

Esta é a novidade chave que distingue o NimbleReg de métodos anteriores baseados em superfície.

1. Construção de Campo de Velocidade Estacionário (SVF): As velocidades esparsas produzidas por $f_\theta$ são convertidas em um SVF ($V$) suave e contínuo usando convolução de kernel.
2. Quadro Log-Euclidiano: O SVF é integrado ao longo do tempo usando a estrutura de grupo de Lie das difeomorfismos. A transformação final $\phi$ é calculada como o mapa exponencial de Lie: $\phi = \exp(V)$.
3. Fusão: Como o quadro SVF permite operações lineares na álgebra de Lie, velocidades de múltiplas regiões anatômicas podem ser mescladas perfeitamente. A transformação final é a composição do pré-alinhamento, o difeomorfismo aprendido e o inverso do pré-alinhamento.
4. Inferência: A transformação resultante pode ser aplicada a qualquer ponto de consulta (por exemplo, uma grade de voxel) ou imagem via mapeamento de coordenadas reverso.

3. Contribuições Principais

• Arquitetura Leve: Ao substituir grades 3D densas por processamento de superfície baseado em PointNet, o modelo reduz drasticamente a pegada de memória e a carga computacional.
• Garantia Difeomórfica via SVF: O artigo conecta com sucesso a lacuna entre registro baseado em superfície e difeomorfismo global integrando o quadro Log-Euclidiano e Campos de Velocidade Estacionários. Isso garante que a transformação final seja suave, invertível e que preserve a topologia.
• Fusão Multi-Região: O método fornece um mecanismo para fundir alinhamentos regionais independentes em uma única transformação global coerente, resolvendo uma limitação de abordagens anteriores de DL baseadas em superfície.
• Eficiência: A abordagem permite inferência rápida com uso mínimo de memória em comparação com métodos de DL baseados em grade de última geração.

4. Resultados Experimentais

O método foi avaliado em uma tarefa de registro cerebral não linear usando dados de MRI ponderados em T1 dos conjuntos de dados ADNI, UK Biobank e IXI (95 regiões anatômicas).

• Comparação: O NimbleReg foi comparado contra o EasyReg (um método de DL de última geração baseado em SynthMorph) e baselines afins/poliafinos tradicionais.
• Métricas de Precisão:
  • Índice Jaccard (Sobreposição de Volume): O NimbleReg alcançou desempenho comparável ao EasyReg em regiões corticais, subcorticais e ventriculares.
  • Distância de Chamfer (Alinhamento de Superfície): O NimbleReg superou o EasyReg em métricas baseadas em superfície. Os autores notam que isso é parcialmente porque o NimbleReg é explicitamente treinado no alinhamento de superfície, enquanto o EasyReg não é.
• Eficiência Computacional:
  • Uso de Memória: O EasyReg requer aproximadamente 24 GB de RAM para um tamanho de lote de 1 devido a mapas de características pesados de CNN. O NimbleReg requer <1 GB para a fase de estimativa de velocidade, sendo a etapa de integração a única parte intensiva em memória (mitigada por KD-trees).
  • Velocidade: A base PointNet permite estimativa de velocidade extremamente rápida.

5. Significado e Trabalho Futuro

Significado:
O NimbleReg demonstra que o registro de imagens difeomórfico de alta qualidade não requer modelos pesados de aprendizagem profunda baseados em grade. Ao aproveitar as propriedades geométricas de superfícies segmentadas e o rigor matemático do quadro Log-Euclidiano, oferece uma alternativa escalável e eficiente em memória, adequada para ambientes com recursos limitados ou estudos em grande escala.

Limitações e Direções Futuras:

• Resolução Cortical: A precisão em regiões corticais poderia ser melhorada com maiores densidades de pontos, pois a decimação atual pode perder detalhes finos.
• Largura de Banda do Kernel: A largura de banda fixa do kernel de integração pode levar a transformações excessivamente suaves em áreas esparsas; um $\sigma$ adaptativo baseado na densidade de pontos é sugerido.
• Esquema de Integração: A integração atual de Runge-Kutta de primeira ordem (Euler) poderia ser atualizada para métodos de ordem superior para melhor precisão numérica.

Em conclusão, o NimbleReg representa um passo significativo em direção a soluções de aprendizagem profunda práticas, eficientes e matematicamente robustas para registro de imagens médicas.