MM-DLS

Multi-task deep learning based on PET/CT images for the diagnosis and prognosis prediction of advanced non-small cell lung cancer

Overview

• Lung-lesion segmentation with cross-attention transformer
• Multi-modal feature fusion (CT, PET, Clinical)
• Multi-task learning: Pathological classification, TNM staging, DFS and OS survival prediction
• Cox proportional hazards survival loss

Key Features

• Multi-modal fusion: Combines CT-based imaging features, PET metabolic biomarkers (SUVmax, SUVmean, SUVpeak, TLG, MTV), and structured clinical variables (age, sex, smoking status, smoking duration, smoking cessation history, tumor size).
• Multi-task learning: Simultaneous optimization for:
  • Histological subtype classification (LUAD vs LUSC)
  • TNM stage classification (I-II, III, IV)
  • Disease-free survival (DFS) prediction
  • Overall survival (OS) prediction
• Attention-based feature fusion: Transformer cross-attention module to integrate lung-lesion spatial information.
• Survival modeling: Incorporates Cox Proportional Hazards loss for survival time prediction.
• Flexible data simulation and loading: Includes utilities for synthetic data generation and multi-slice 2D volume processing.

Architecture

1. Segmentation Module (LungLesionSegmentor):
   • Shared ResNet encoder to extract features from CT images.
   • Dual decoders for lung and lesion segmentation.
   • Transformer-based cross-attention module for enhanced spatial feature interaction between lung and lesion regions.
2. Feature Encoders:
   • LesionEncoder: 2D convolutional encoder for lesion patches.
   • SpaceEncoder: 2D convolutional encoder for lung-space contextual patches.
3. Attention Fusion Module:
   • LesionAttentionFusion: Multi-head attention to fuse lesion and lung features into compact patient-level representations.
4. Patient-Level Fusion Model (PatientLevelFusionModel):
   • Fully connected network that combines imaging, PET, and clinical features.
   • Outputs classification logits, DFS and OS risk scores.
5. Loss Functions:
   • Binary cross-entropy loss for classification.
   • Cox proportional hazards loss (CoxPHLoss) for survival prediction.

Code Structure

• ModelLesionEncoder.py: Lesion image encoder extracting discriminative features from multi-slice tumor regions.
• ModelSpaceEncoder.py: Lung space encoder modeling anatomical and spatial context beyond the lesion.
• LesionAttentionFusion.py: Attention-based fusion module for adaptive integration of lesion and spatial features.
• ClinicalFusionModel.py: Patient-level fusion network combining imaging features, radiomics, PET signals, and clinical variables.
• HierMM_DLS.py：Core hierarchical multimodal deep learning model supporting multi-task learning: (1)Subtype classification; (2)TNM stage prediction; (3)DFS and OS modeling
• CoxphLoss.py: Cox proportional hazards loss for survival modeling with censored data.
• PatientDataset.py:Patient dataset loader supporting imaging, radiomics, PET, clinical variables, survival outcomes, and treatment labels.
• LungLesionSegmentation.py: Lung-lesion segmentation model
• ImageDataLoader.py: Image preprocessing and loading utilities for multi-slice inputs.
• plot_results.py: Visualization utilities for Kaplan–Meier curves, hazard ratios, and survival analysis results.

Data Format

Imaging Data:

• CT slices (PNG format)
• Lung masks (binary masks, PNG)
• Lesion masks (binary masks, PNG)
• Slices grouped per patient ID

Tabular Data:

• Radiomics features: 128-dimensional vector (PyRadiomics extracted)
• PET features: [SUVmax, SUVmean, SUVpeak, TLG, MTV]
• Clinical features: [Age, Sex, Smoking Status, Smoking Duration, Smoking Cessation, Tumor Diameter]
• Survival data: DFS time/event, OS time/event
• Classification label: LUAD (0) or LUSC (1)

Installation

Install dependencies

Usage

🔽 Download Pretrained Models

• MM-DLS (Full multimodal, best checkpoint)
  ⬇️ Download Pretrained Model Size 1.3 MB The MM-DLS model is intentionally lightweight (~1.3 MB), as it employs compact CNN encoders and MLP-based multimodal fusion rather than large pretrained backbones, enabling efficient deployment and fast inference.

Model Performance (from publication)

Histological Subtype Classification:

TNM Stage Prediction:

AP and calibration maintained across internal and external sets

Reference