Awesome

PA-Diff

The code of paper :

Learning A Physical-aware Diffusion Model Based on Transformer for Underwater Image Enhancement

Chen Zhao, Chenyu Dong, Weiling Cai

You can find our paper on Arxiv

Dependencies

• PyTorch
• einops
• opencv-python

Installation

You should install Pytorch first (This installation command may not be available, please find the appropriate command in PyTorch )

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

Use requirements.txt to install other requirements

pip install -r requirements.txt

Datasets

LSUI｜UIEB｜U45｜Realworld Images

Prepare dataset:

1. You should make sure your image size is 256*256
2. Your directory structure should look like this:

train:
DatasetName_train_16_256  (16 and 256 are set in config)
├── hr_256 (GT images)
└── sr_16_256 (input images)

test:
DatasetName_val_16_256 
├── hr_256
└── sr_16_256

Test

1. Download our checkpoints from Baidu Netdisk(We will upload our checkpoint after our paper is accepted.)
2. Change "resume_state" in the config file to the path of your checkpoints
3. Change test "dataroot" in the config file to the path like xxx/DatasetName_val_16_256
4. Run : python infer.py

Train

1. Change train "dataroot" in the config file to the path like xxx/DatasetName_train_16_256
2. Change "resume_state" in the config file to "null"
3. Run : python train.py

Citation

@article{zhao2024learning,
  title={Learning A Physical-aware Diffusion Model Based on Transformer for Underwater Image Enhancement},
  author={Zhao, Chen and Dong, Chenyu and Cai, Weiling},
  journal={arXiv preprint arXiv:2403.01497},
  year={2024}
}