--- license: mit language: - en --- # 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 **MM-DLS** is a multi-modal, multi-task deep learning framework for the diagnosis, staging, and prognosis prediction of advanced non-small cell lung cancer (NSCLC). It integrates multi-source data including CT images, PET metabolic parameters, and clinical information to provide a unified, non-invasive decision-making tool for personalized treatment planning. This repository implements the full MM-DLS pipeline, consisting of: - 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 The framework supports both classification (adenocarcinoma vs squamous cell carcinoma) and survival risk prediction tasks, and has been validated on large-scale multi-center clinical datasets. --- ## 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 The overall MM-DLS system consists of: ![Python](https://img.shields.io/badge/python-3.9%2B-blue) ![PyTorch](https://img.shields.io/badge/PyTorch-2.x-red) ![CUDA](https://img.shields.io/badge/CUDA-11.8%2B-green) ![License](https://img.shields.io/badge/License-MIT-lightgrey) ![Status](https://img.shields.io/badge/Status-Research-orange) 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 The input data is organized per patient as follows: ### 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) Simulated data utilities are provided for experimentation and reproducibility. --- ## Installation ```bash # Clone repository conda create -n mm_dls python=3.10 -y conda activate mm_dls git clone https://github.com/your_username/MM-DLS-NSCLC.git ``` ## Install dependencies ```bash pip install -r requirements.txt ``` ## Usage ### 🔽 Download Pretrained Models Pretrained MM-DLS models are available for direct download: - **MM-DLS (Full multimodal, best checkpoint)** [⬇️ Download Pretrained Model](https://drive.google.com/file/d/1IcyCwMgCX8wv0NMp84U4wlzhLoXH7ayx/view?usp=drive_link) 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. After downloading, place the model files under the `./MODEL/` directory: Training: ```bash python train_patient_model.py ``` Evaluation: ```bash python test.py ``` Example Forward Pass: ```bash python run_sample.ipynb ``` ## Model Performance (from publication) ### Histological Subtype Classification: AUC: 0.85 ~ 0.92 across cohorts AP: 0.81 ~ 0.86 ### TNM Stage Prediction: AUC: Stage I-II (0.86 ~ 0.96), Stage III (0.85 ~ 0.95), Stage IV (0.83 ~ 0.95) ### AP and calibration maintained across internal and external sets DFS & OS Prognosis: C-index: up to 0.75 Time-dependent AUC (1/2/3 years): 0.77 ~ 0.91 Brier score: consistently < 0.2 for DFS and < 0.3 for OS Superior to single modality models (clinical-only or imaging-only) ## Reference Please cite our original publication when using this work: License This project is licensed under the MIT License. ⚠️ **Notice:** The pretrained model is shared solely for research validation purposes and **should not be used, distributed, or cited before the associated study is formally published**. Contact For any questions or collaborations, please contact: Dr. Fang Dai: daifang_cool@163.com