PyTorch-StratifyMEC

A Multimodal Deep Learning Framework with Contrastive Learning and Multi-Instance Learning for Endometrial Cancer Preoperative Risk Stratification

Overview

PyTorch-StratifyMEC is a robust multimodal deep learning framework designed for the preoperative risk stratification of endometrial cancer (low-risk vs. non-low-risk). By seamlessly integrating multiparametric MRI (mpMRI) imaging and MRI radiology reports, this framework leverages advanced deep learning techniques to enhance diagnostic accuracy across multicenter and international cohorts.

Core Innovations:

• Multimodal Fusion: Combines visual features from MRI and semantic representations from medical text using a hybrid fusion strategy (early feature-level + late prediction-level).
• Report-Guided Contrastive Learning: Aligns visual and textual representations to bridge the semantic gap between imaging findings and radiologist interpretations.
• Multi-Instance Learning (MIL): Effectively aggregates slice-level features into patient-level risk predictions, enabling flexible handling of variable MRI sequence availability.
• Cross-Attention Mechanism: Utilizes bidirectional cross-attention and feature re-embedding to capture fine-grained correlations between modalities at multiple interaction stages.

Framework Architecture

The framework operates through three primary stages:

1. Segmentation Module: Employs nnU-Net for volumetric tumor segmentation. For each available MRI sequence (T1CA, T1CS, T2A, T2S), the slice with the maximum tumor area is selected for downstream processing. The segmentation is performed in 3D using the official nnU-Net v2 framework.

2. Encoding Stage:

   • Text Encoder: A BERT-based architecture initialized with MedBERT (publicly available Chinese medical pretrained weights), utilizing a 768-dimensional pooled representation with a dropout rate of 0.85.
   • Image Encoder: ResNet50 pre-trained on ImageNet for robust visual feature extraction.

3. Fusion & Prediction: Integrates image-text contrastive learning, bidirectional cross-attention, feature re-embedding, and MIL aggregation. The model combines early feature-level fusion with late prediction-level fusion to maximize risk stratification performance. No manually predefined weighting coefficients are assigned; instead, modality contributions are learned jointly through cross-attention, MIL aggregation, and final fusion.

Install

Set up the environment using Conda and install the required dependencies:

conda create -n StratifyMEC python=3.12
conda activate StratifyMEC
pip install -r requirements.txt

Data Preparation

The dataset directory format is as follows:

── texts.xlsx
├── patient1
│   ├── xxxx1.jpg or png
│   ├── xxxx2.jpg or png
│   ├── xxxx3.jpg or png
│   ── xxxx4.jpg or png
├── patient2
│   ├── xxxx1.jpg or png
│   ├── xxxx2.jpg or png
│   ├── xxxx3.jpg or png
│   ...

The content of texts.xlsx is as follows:

text	label	path	center
Imaging Findings:xxx.Report Conclusion:xxx.	0	patient1	center1
Imaging Findings:xxx.Report Conclusion:xxx.	1	patient2	center1
...	...	...	...

• text: MRI report text, with "Imaging Findings" and "Report Conclusion" sections concatenated
• label: 0 = low-risk, 1 = non-low-risk
• path: Patient-level image directory path (without file extension)
• center: Source center identifier

International Cohort Data (Publicly Available)

To facilitate external validation and reproducibility, we have publicly released the international evaluation cohort data.

• Data Sources: CPTAC-UCEC (The Cancer Proteome Atlas - Uterine Corpus Endometrial Carcinoma) and TCGA-UCEC (The Cancer Genome Atlas - Uterine Corpus Endometrial Carcinoma).
• Geographic Origin: USA.
• Sample Size: 12 cases (6 from CPTAC-UCEC + 6 from TCGA-UCEC).
• Risk Distribution: 3 Low-Risk, 9 Non-Low-Risk.
• Purpose: Independent international validation set to assess the cross-population generalizability of the StratifyMEC framework.

The international cohort data is already included in this repository.

Pre-trained Weights (Publicly Available)

To support downstream research and clinical translation, we have publicly released the following pre-trained model weights:

Model Component	Description
Text Encoder	BERT weights initialized with MedBERT (Chinese medical pretrained weights), fine-tuned on MRI reports
Image Encoder	ResNet50 weights with image-text contrastive learning pre-training
Fusion Model	Complete multimodal fusion network weights (early + late fusion)

The pre-trained weights is available for download at: OneDrive Link .

Training

Configure the parameters in the respective configuration files and execute the following training scripts in sequence:

Step 1: Train Text Encoder (BERT)

python train_bert.py

Step 2: Train Image Encoder with Contrastive Learning

python train_clip_model.py

Step 3: Train Fusion Network

python train_fusion.py

Evaluation

Evaluate the trained model on validation and test cohorts:

python test_fusion.py

Citation

If you use this code or data in your research, please cite our paper:

Kang H, Wang Y, He H, Cai G, Zou Y, Li J, Shen Y, Zhuo J, Aili A, Yang L, Yang L, Zheng W, Zhang L, Zhuang P, Lin D, Ban X, Duan X, Cheng S. A Multimodal Deep Learning Framework with Contrastive Learning and Multi-Instance Learning for Endometrial Cancer Preoperative Risk Stratification.