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} }