XNet

XNet

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

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:

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:

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:

python train.py --dataset_name YourDatasetName

Disable Components (Ablation Study)

# 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

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.