Update README.md

README.md

@@ -2,4 +2,176 @@
 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