WaveDE: Wavelet-Routed Dual-Expert Network for Fundus Photograph Enhancement

Status: This work has been provisionally accepted for presentation at MICCAI 2026.

This repository is the official implementation of the paper "WaveDE: Wavelet-Routed Dual-Expert Network for Fundus Photograph Enhancement."

Authors

Haitao Nie^1,†, Dongjie Wu^2,†, Bin Xie^1,✉

^† These authors contributed equally (co-first authors). ^✉ Corresponding author.

Affiliations

¹ School of Automation, Central South University, Changsha, Hunan, China

² School of Computer Science and Engineering, Central South University, Changsha, Hunan, China

---

Getting Started

Repository Structure

WaveDE/
├── WaveDE.py                  # Training entry point
├── infer_real_image.py        # Inference on real fundus images
├── datasets.py                # Dataset loading and preprocessing
├── loss.py                    # Loss functions (Charbonnier, MS-SSIM, FFL, and ACR)
├── moe_modules.py             # FD-Block, LF-VMamba, OS-HFNet, and HFI modules
├── skmoe.py                   # WaveDE assembly (HFI Enhanced VAE Decoder + diffusion U-Net with FD-Blocks)
├── utils.py                   # Utility functions (EMA, logging, etc.)
├── segment/
│   ├── PraNet_Res2Net.py      # PraNet-based optic-disc teacher network
│   └── Res2Net_v1b.py         # Res2Net backbone
└── README.md

---

Data Naming Convention

PairFolder (in datasets.py) matches LQ, HQ by filename stem. The rules are:

Data	Naming Rule	Example
HQ	Any name; its stem is the matching key	fundus.png
LQ	Must end with _<digits>; the part before _digits must match an HQ stem	fundus_0.png → matches HQ fundus.png

Supported formats: .png, .jpg, .jpeg, .bmp, .tif, .tiff (case-insensitive).

One HQ image can pair with multiple LQ images (one-to-many). Directories are scanned recursively.

Example directory layout:

hq/
  fundus.png
  retina.png

lq/
  fundus_0.png        # matches hq/fundus.png
  fundus_1.png        # matches hq/fundus.png
  retina_0.png        # matches hq/retina.png

Note: LQ files that do not end with _<digits> are silently skipped. If no valid pairs are found, training raises a RuntimeError.

---

Environment

# 1. Create a conda environment
conda create -n WaveDE python=3.10 -y
conda activate WaveDE

# 2. Install PyTorch
# Example for CUDA 12.1 (replace to match your CUDA version):
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

# 3. Install other dependencies
pip install numpy tqdm Pillow diffusers einops pytorch_wavelets PyWavelets piq lpips segmentation-models-pytorch swanlab tensorboard matplotlib

# 4. Install mamba-ssm
pip install causal-conv1d --no-build-isolation
pip install mamba-ssm --no-build-isolation

---

Training

python WaveDE.py \
  --hq_root /path/to/hq/train \
  --lq_root /path/to/degradation/train/lq \
  --out /path/to/output \
  --test_hq_root /path/to/hq/validation \
  --test_lq_root /path/to/degradation/validation/lq \
  --model_id stable-diffusion-v1-5/stable-diffusion-v1-5 \
  --path_vessel /path/to/vessel/best_model.pth \
  --path_od /path/to/PraNet-Best.pth \
  --steps 30000 \
  --batch 32 \
  --accum 1 \
  --lr 1e-4 \
  --warmup_steps 1000 \
  --timestep 200 \
  --cfp_l1 1.0 \
  --cfp_msssim 0.25 \
  --lambda_vessel 0.01 \
  --lambda_od 0.04 \
  --lambda_ffl 0.1 \
  --ema_decay 0.999 \
  --val_every 4000 \
  --grad_clip_norm 1.0 \
  --bf16 \
  --gradient_checkpointing \
  --compile \
  --ema

To resume training, append --resume_from /path/to/checkpoint.pth.

---

Inference

python infer_real_image.py \
  --weights /path/to/best_model.pth \
  --model_id stable-diffusion-v1-5/stable-diffusion-v1-5 \
  --out_dir /path/to/output \
  --test_lq_root /path/to/test/lq \
  --res 256 \
  --timestep 200 \
  --batch 32

---

License

License: CC BY-NC-SA 4.0

You are free to share and adapt the contents of this repository under the following terms:

• Attribution (BY)
• NonCommercial (NC)
• ShareAlike (SA)

Full text: https://creativecommons.org/licenses/by-nc-sa/4.0/

---

Citation

To be updated.

% Citation information will be added here (To be updated).