FEUNet-PyTorch

This is a PyTorch implementation for FEUNet image denoising. Paper download: FEUNet

How to run

1. Dependences

• PyTorch
• torchvision
• OpenCV
• HDF5 for Python
• tensorboardX (TensorBoard for PyTorch)

2. Data Preparation

2.1 dataset download

• The Waterloo Exploration Database

2.2 color image

    python prepare_patches.py --p 50

2.3 grayscale image

    python prepare_patches.py --p 64 --gray

NOTES --max_number_patches can be used to set the maximum number of patches contained in the database. --aug_times can be used to set the number of data augmentation, we set it as 5.

3. Train FEUNet

3.1 grayscale image

python train_gray.py --batch_size 128 --epochs 80 --noiseIntL 0 75 --val_noiseL 25

3.2 color image

python train_rgb.py --batch_size 128 --epochs 80 --noiseIntL 0 75 --val_noiseL 25

4. Test FEUNet

4.1 grayscale image

python avg_gray_test.py --input Set12 --noise_sigma 25 --add_noise True

To run the algorithm on CPU instead of GPU:

python avg_gray_test.py --input Set12 --noise_sigma 25 --add_noise True --no_gpu

4.2 color image

python avg_rgb_test.py --input Kodak24 --noise_sigma 25 --add_noise True

To run the algorithm on CPU instead of GPU:

python avg_rgb_test.py --input Kodak24 --noise_sigma 25 --add_noise True --no_gpu

5. Network Structure

6. Results

6.1 gaussian grayscale image denoising

6.2 gaussian color image denoising

6.3 real image denoising

6.4 image smoothing

6.5 model computational complexity comparison

7. Citation

If you find FEUNet helpful for your research, please consider citing our work.

@article{Wu2023,
  author       = {Wencong Wu and
                  Guannan Lv and
                  Shicheng Liao and
                  Yungang Zhang},
  title        = {FEUNet: a flexible and effective U-shaped network for image denoising},
  journal      = {Signal Image Video Process.},
  volume       = {17},
  number       = {5},
  pages        = {2545--2553},
  year         = {2023}
}