MTC-HSDNet: Multi-Task Collaboration and Hierarchical Self-Distillation in a Multi-Level Transformer Fusion Network for Wilms' Tumor Segmentation and Metastasis Prediction

This repository is the official implementation of MTC-HSDNet.

A deep learning framework for simultaneous Wilms' tumor segmentation and metastasis prediction using multi-task collaboration and hierarchical self-distillation in a multi-level transformer fusion network.

📄 Framework

---

Overview

MTC-HSDNet is a specialized deep learning framework designed for Wilms' tumor analysis, combining multi-task collaboration and hierarchical self-distillation for simultaneous tumor segmentation and metastasis prediction. The network architecture incorporates:

• Multi-Level Transformer Fusion for comprehensive feature integration
• Hierarchical Self-Distillation for progressive knowledge transfer
• Multi-Task Collaboration between segmentation and classification tasks
• Swin Transformer backbone for robust feature extraction
• Feature Pyramid Network (FPN) for multi-scale feature fusion
• Mixture of Experts (MoE) for adaptive feature processing

Requirements

See requirements.txt for detailed dependencies. Key requirements include:

• Python 3.8+
• PyTorch 2.4.0+
• MONAI 1.4.0+
• NumPy, SciPy, scikit-learn
• SimpleITK for medical image processing

Installation

1. Clone the repository:

git clone git@github.com:wenjing-gg/MTC-HSDNet.git
cd MTC-HSDNet

2. Install dependencies:

pip install -r requirements.txt

Usage

Training

python train.py --data_path /path/to/your/data --epochs 100 --batch_size 2

Key Parameters

• --data_path: Path to the dataset directory
• --weights: Path to pre-trained weights (optional)
• --epochs: Number of training epochs
• --batch_size: Batch size for training
• --freeze_layers: Whether to freeze backbone layers

Project Structure

MTC-HSDNet/
├── models/                 # Core model implementations
│   ├── mtc_hsdnet.py      # Main MTC-HSDNet architecture
│   ├── FPN.py             # Feature Pyramid Network
│   ├── kan.py             # Kolmogorov-Arnold Network components
│   ├── loss.py            # Multi-task loss functions
│   └── __init__.py
├── data/                   # Data loading and preprocessing
│   ├── dataset.py         # Custom NRRD dataset implementation
│   └── __init__.py
├── utils/                  # Utility functions
│   ├── utils.py           # Training and evaluation utilities
│   ├── metrics.py         # Evaluation metrics
│   └── __init__.py
├── train.py               # Training script
├── requirements.txt       # Dependencies
└── README.md

Dataset Structure

The expected dataset structure:

your_data/
├── train/
│   ├── 0/  # No metastasis
│   │   ├── image1.nrrd
│   │   ├── image1_label.nrrd
│   │   └── ...
│   └── 1/  # Metastasis
│       ├── image2.nrrd
│       ├── image2_label.nrrd
│       └── ...
└── test/
    ├── 0/
    └── 1/

Citation

If you use this code in your research, please cite:

@article{mtc-hsdnet,
  title={MTC-HSDNet: Multi-Task Collaboration and Hierarchical Self-Distillation in a Multi-Level Transformer Fusion Network for Wilms' Tumor Segmentation and Metastasis Prediction},
  author={Your Name},
  journal={Your Journal},
  year={2024}
}

License

This project is licensed under the MIT License - see the LICENSE file for details.

Contact

For questions or issues, please contact: [yuwenjing259@gmail.com]