|
@@ -1,3 +1,166 @@
|
|
|
# XNet
|
|
# XNet
|
|
|
|
|
|
|
|
-XNet: A Staged Dual-Frequency Synergistic Framework via Wavelet-FFT for Medical Image Segmentation of Small Objects and Weak Boundaries
|
|
|
|
|
|
|
+XNet: A Staged Dual-Frequency Synergistic Framework via Wavelet-FFT for Medical Image Segmentation of Small Objects and Weak Boundaries
|
|
|
|
|
+
|
|
|
|
|
+## Overview
|
|
|
|
|
+
|
|
|
|
|
+XNet is a medical image segmentation framework designed for accurate segmentation of small objects and weak boundaries. It combines wavelet transform and FFT enhancement with Swin-UNETR architecture to achieve superior performance in challenging scenarios.
|
|
|
|
|
+
|
|
|
|
|
+## Key Features
|
|
|
|
|
+
|
|
|
|
|
+- **Dual-Frequency Enhancement**: Combines wavelet transform and FFT for multiscale feature extraction
|
|
|
|
|
+- **Swin-UNETR Backbone**: Leverages transformer-based architecture for robust segmentation
|
|
|
|
|
+- **Enhanced Data Augmentation**: Comprehensive augmentation pipeline for better generalization
|
|
|
|
|
+- **Multi-Metric Evaluation**: Dice, IoU, and Hausdorff Distance metrics
|
|
|
|
|
+
|
|
|
|
|
+## Installation
|
|
|
|
|
+
|
|
|
|
|
+### Prerequisites
|
|
|
|
|
+
|
|
|
|
|
+```bash
|
|
|
|
|
+pip install torch monai swanlab opencv-python numpy
|
|
|
|
|
+```
|
|
|
|
|
+
|
|
|
|
|
+## Quick Start
|
|
|
|
|
+
|
|
|
|
|
+### 1. Prepare Your Dataset
|
|
|
|
|
+
|
|
|
|
|
+Organize your dataset in the following structure:
|
|
|
|
|
+
|
|
|
|
|
+```
|
|
|
|
|
+data/
|
|
|
|
|
+└── Polyp-Detection-Dataset/
|
|
|
|
|
+ └── YourDatasetName/
|
|
|
|
|
+ ├── images/
|
|
|
|
|
+ ├─────1.png
|
|
|
|
|
+ ├─────2.png
|
|
|
|
|
+ ├─────...
|
|
|
|
|
+ └── masks/
|
|
|
|
|
+ ├─────1.png
|
|
|
|
|
+ ├─────2.png
|
|
|
|
|
+ ├─────...
|
|
|
|
|
+ └── train.txt
|
|
|
|
|
+ └── val.txt
|
|
|
|
|
+```
|
|
|
|
|
+
|
|
|
|
|
+### 2. Training
|
|
|
|
|
+
|
|
|
|
|
+Train the model on your dataset:
|
|
|
|
|
+
|
|
|
|
|
+```bash
|
|
|
|
|
+python train.py \
|
|
|
|
|
+ --dataset_name YourDatasetName \
|
|
|
|
|
+ --data_root ./data/Polyp-Detection-Dataset \
|
|
|
|
|
+ --batch_size 4 \
|
|
|
|
|
+ --max_epochs 1000 \
|
|
|
|
|
+ --learning_rate 1e-4 \
|
|
|
|
|
+ --device cuda
|
|
|
|
|
+```
|
|
|
|
|
+
|
|
|
|
|
+#### Key Training Parameters
|
|
|
|
|
+
|
|
|
|
|
+| Parameter | Default | Description |
|
|
|
|
|
+|-----------------------------|------------|----------------------------|
|
|
|
|
|
+| `--dataset_name` | Required | Name of your dataset |
|
|
|
|
|
+| `--batch_size` | 4 | Batch size for training |
|
|
|
|
|
+| `--max_epochs` | 1000 | Maximum training epochs |
|
|
|
|
|
+| `--learning_rate` | 1e-4 | Initial learning rate |
|
|
|
|
|
+| `--feature_size` | 48 | Network feature dimension |
|
|
|
|
|
+| `--target_spatial_size` | (512, 512) | Input image size |
|
|
|
|
|
+| `--early_stopping_patience` | 100 | Early stopping patience |
|
|
|
|
|
+| `--use_wavelet` | True | Enable wavelet enhancement |
|
|
|
|
|
+| `--use_fft` | True | Enable FFT enhancement |
|
|
|
|
|
+
|
|
|
|
|
+### 3. Evaluation
|
|
|
|
|
+
|
|
|
|
|
+Evaluate trained models:
|
|
|
|
|
+
|
|
|
|
|
+```bash
|
|
|
|
|
+python eval.py \
|
|
|
|
|
+ --dataset_name YourDatasetName \
|
|
|
|
|
+ --data_root ./data/Polyp-Detection-Dataset \
|
|
|
|
|
+ --outputs_dir ./outputs \
|
|
|
|
|
+ --device cuda
|
|
|
|
|
+```
|
|
|
|
|
+
|
|
|
|
|
+#### Key Evaluation Parameters
|
|
|
|
|
+
|
|
|
|
|
+| Parameter | Default | Description |
|
|
|
|
|
+|------------------------|------------------|---------------------------------------------|
|
|
|
|
|
+| `--dataset_name` | Required | Name of your dataset |
|
|
|
|
|
+| `--outputs_dir` | ./outputs_minute | Directory containing trained models |
|
|
|
|
|
+| `--batch_size` | 1 | Batch size for evaluation |
|
|
|
|
|
+| `--save_visualization` | True | Save visualization results |
|
|
|
|
|
+| `--vis_num_samples` | 1000 | Number of samples to visualize |
|
|
|
|
|
+| `--best_metric` | False | Use best overall model (default: best Dice) |
|
|
|
|
|
+
|
|
|
|
|
+## Model Architecture
|
|
|
|
|
+
|
|
|
|
|
+XNet integrates three key components:
|
|
|
|
|
+
|
|
|
|
|
+1. **Wavelet Enhancement Module**: Captures multi-scale frequency features
|
|
|
|
|
+2. **FFT Enhancement Module**: Enhances global frequency domain information
|
|
|
|
|
+3. **Swin-UNETR v2**: Transformer-based backbone for robust feature extraction
|
|
|
|
|
+
|
|
|
|
|
+## Output Structure
|
|
|
|
|
+
|
|
|
|
|
+After training, outputs are organized as:
|
|
|
|
|
+
|
|
|
|
|
+```
|
|
|
|
|
+outputs_minute/
|
|
|
|
|
+├── best_dice_model_YourDatasetName.pt
|
|
|
|
|
+├── best_iou_model_YourDatasetName.pt
|
|
|
|
|
+├── best_metric_model_YourDatasetName.pt
|
|
|
|
|
+└── checkpoints_YourDatasetName/
|
|
|
|
|
+ └── checkpoint_epoch=X.pt
|
|
|
|
|
+```
|
|
|
|
|
+
|
|
|
|
|
+## Monitoring with SwanLab
|
|
|
|
|
+
|
|
|
|
|
+Training progress is automatically logged to SwanLab:
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+Metrics tracked:
|
|
|
|
|
+- Training/validation loss
|
|
|
|
|
+- Dice coefficient
|
|
|
|
|
+- IoU (Intersection over Union)
|
|
|
|
|
+- Hausdorff Distance
|
|
|
|
|
+- Learning rate schedule
|
|
|
|
|
+
|
|
|
|
|
+## Advanced Usage
|
|
|
|
|
+
|
|
|
|
|
+### Resume Training
|
|
|
|
|
+
|
|
|
|
|
+Training automatically resumes from the latest checkpoint:
|
|
|
|
|
+
|
|
|
|
|
+```bash
|
|
|
|
|
+python train.py --dataset_name YourDatasetName
|
|
|
|
|
+```
|
|
|
|
|
+
|
|
|
|
|
+### Disable Components (Ablation Study)
|
|
|
|
|
+
|
|
|
|
|
+```bash
|
|
|
|
|
+# Disable wavelet enhancement
|
|
|
|
|
+python train.py --dataset_name YourDatasetName --no_wavelet
|
|
|
|
|
+
|
|
|
|
|
+# Disable FFT enhancement
|
|
|
|
|
+python train.py --dataset_name YourDatasetName --no_fft
|
|
|
|
|
+```
|
|
|
|
|
+
|
|
|
|
|
+### Custom Loss Weights
|
|
|
|
|
+
|
|
|
|
|
+```bash
|
|
|
|
|
+python train.py \
|
|
|
|
|
+ --dataset_name YourDatasetName \
|
|
|
|
|
+ --dice_weight 1.0 \
|
|
|
|
|
+ --ce_weight 1.0 \
|
|
|
|
|
+ --iou_weight 1.0
|
|
|
|
|
+```
|
|
|
|
|
+
|
|
|
|
|
+## Citation
|
|
|
|
|
+
|
|
|
|
|
+If you find this work useful, please cite our paper.
|
|
|
|
|
+
|
|
|
|
|
+## License
|
|
|
|
|
+
|
|
|
|
|
+This project is licensed under the Apache License.
|