Classificação de Nódulos Pulmonares com VMamba e ResNet50

Este repositório contém a implementação oficial do código utilizado para a monografia sobre classificação binária de nódulos pulmonares (Benigno vs. Maligno) em imagens de Tomografia Computadorizada (TC).

O estudo compara o desempenho de uma arquitetura baseada em State Space Models (VMamba/SS2D) contra uma baseline convolucional (ResNet50), utilizando uma abordagem de pré-processamento 2.5D.

🗂️ Base de Dados

O projeto utiliza dados públicos provenientes de duas fontes principais, cruzados para obter anotações precisas de malignidade:

• LUNA16 (LUng Nodule Analysis 2016): Fornece as coordenadas (X, Y, Z) e o diâmetro dos nódulos candidatos, além das imagens brutas (.mhd/.raw).
• LIDC-IDRI: Fornece as classificações de malignidade (escala 1 a 5) atribuídas por radiologistas.

Preparação dos Dados (data_preparation.py)

O script realiza o download automático dos subsets do LUNA16 via Zenodo e executa o cruzamento lógico com o LIDC:

Nota: O script espera que o arquivo Filtered_lidc_malignancy_mean.csv (derivado do LIDC-IDRI) seja manualmente carregado para o diretório data/.

1. Matching: Nódulos são pareados baseados na distância euclidiana e similaridade de diâmetro.
2. Rotulagem:
   • Benigno (Classe 0): Score de malignidade médio $\le 2.5$.
   • Maligno (Classe 1): Score de malignidade médio $\ge 3.5$.
   • Nódulos indeterminados (entre 2.5 e 3.5) são descartados para reduzir ruído de anotação.

⚙️ Pré-processamento e Extração de Features

O pipeline de pré-processamento (pre_processing.py) converte os volumes 3D de TC em representações 2.5D compatíveis com as arquiteturas de visão computacional.

• Conversão de Coordenadas: Transformação de coordenadas do "Mundo" (mm) para coordenadas de "Voxel" baseadas na origem e espaçamento da imagem .mhd.
• Extração 2.5D: Para cada nódulo, são extraídos 3 cortes axiais adjacentes ($z-1, z, z+1$). Estes cortes formam os 3 canais da imagem de entrada (RGB-like).
• Janelamento Hounsfield: Aplicação de janela pulmonar rígida:
  • Min: -1000 HU
  • Max: +400 HU
• Normalização e Resize: As imagens são normalizadas para $[0, 1]$ e redimensionadas para 224x224 pixels, salvas em formato .npy.

🧠 Modelos Implementados (models.py)

1. VMamba (Visual Mamba / SS2D)

• Implementação baseada no mecanismo Selective Scan 2D (SS2D).
• Utiliza o Cross-Scan Module para processar a imagem em 4 direções (original, flip horizontal, flip vertical, duplo flip), permitindo que o modelo Mamba capture dependências globais na imagem sem o custo quadrático da Self-Attention.
• Configuração Tiny adaptada para classificação binária.

2. ResNet50 (Baseline)

• Backbone padrão pré-treinado na ImageNet.
• A última camada linear é substituída para gerar um único logit de saída.

🧪 Estratégia de Treinamento e Validação

O protocolo experimental (train.py e data_loaders.py) foi desenhado para garantir robustez, evitar vazamento de dados (data leakage) e lidar com o desbalanceamento de classes.

• Divisão de Dados Robusta:

  • Hold-Out Set: 15% dos pacientes são completamente separados para o teste final.
  • Validação Cruzada (Group K-Fold): Os 85% restantes são divididos em 5 folds. A divisão é baseada no seriesuid (ID do paciente), garantindo que nódulos do mesmo paciente nunca estejam em conjuntos de treino e validação simultaneamente.

• Técnicas para Desbalanceamento:

  • Amostragem Ponderada: No treino, WeightedRandomSampler é utilizado para criar batches com proporção equilibrada de exemplos benignos e malignos.
  • Função de Perda (Focal Loss): BinaryFocalLoss é usada para dar mais peso aos exemplos difíceis e mal classificados, melhorando a convergência em cenários desbalanceados.

• Otimização e Fine-Tuning:

  • Otimizador: AdamW com weight_decay.
  • Learning Rate Scheduler: CosineAnnealingLR com um período de warmup linear no início do treino para estabilizar a convergência, especialmente para o VMamba.
  • Learning Rate Diferencial: O backbone pré-treinado utiliza um learning rate menor (e.g., 1e-6) que a cabeça de classificação (e.g., 1e-4) para preservar as features aprendidas.
  • Treinamento em Duas Fases (Opcional para VMamba): O script train.py suporta uma estratégia onde primeiro se treina apenas a cabeça de classificação com o backbone congelado e, em uma segunda fase, o modelo inteiro é fine-tuned com um learning rate baixo.

• Parada e Seleção de Modelo:

  • Early Stopping: O treino é interrompido se a métrica AUC-ROC no conjunto de validação não melhorar por um número definido de épocas (patience).
  • Seleção de Checkpoint: O melhor checkpoint de cada fold é salvo com base no maior valor de AUC de validação.

📂 Estrutura do Repositório

.
├── data/                       # (Gerado) Armazena .mhd, .raw, CSVs e patches .npy
├── figures_monografia/         # (Gerado) Figuras para o TCC
├── models/                     # (Gerado) Checkpoints dos modelos treinados (.pth)
├── weights/                    # Pesos pré-treinados do VMamba (vssmtiny_...pth)
├── main.ipynb                  # Notebook orquestrador do pipeline
├── data_preparation.py         # Download e cruzamento de bases (LUNA16 + LIDC)
├── pre_processing.py           # Extração de patches 2.5D e conversão HU
├── data_loaders.py             # Dataset PyTorch, K-Fold e Augmentations
├── models.py                   # Definição das arquiteturas VMamba e ResNet50
├── train.py                    # Loop de treino, validação e checkpointing
├── metrics.py                  # Funções para cálculo de métricas
├── figures_for_monography.py   # Scripts para gerar visualizações
└── results_holdout.json        # (Gerado) Resultados agregados da avaliação no hold-out

🚀 Como Executar

A execução é orquestrada através do main.ipynb, idealmente em ambiente Google Colab (devido à integração com GDrive para persistência de dados).

1. Configure as variáveis de ambiente no notebook.
2. Execute a célula de Preparação para baixar e organizar os CSVs.
3. Execute o Pré-processamento para gerar os arquivos .npy (processo demorado, ~60GB de dados).
4. Configure os hiperparâmetros em train.py (ou via chamadas no notebook) e inicie o treinamento.

Exemplo de execução de treino (via código)

from train import run_training

# Treino da ResNet50 no Fold 0
run_training(
    model_name="resnet50",
    fold=0,
    batch_size=32,
    epochs=50
)

📊 Métricas Avaliadas

O desempenho é medido no conjunto de hold-out (15% dos pacientes, isolados antes do K-Fold) utilizando:

• AUC-ROC (Métrica primária)
• Sensibilidade (Recall da classe Maligna)
• Especificidade (Recall da classe Benigna)
• F1-Score
• Acurácia

📄 Licença e Créditos

Este código foi desenvolvido para fins estritamente acadêmicos.

• LUNA16 Dataset: https://zenodo.org/record/2595813
• LIDC-IDRI Dataset: https://wiki.cancerimagingarchive.net/display/Public/LIDC-IDRI