[NCHU thesis]
Image Restoration by Selective Residual Block on Improved Hierarchical Encoder-Decoder Networks

🎉 This thesis received the 2022 Best Dissertation Award from the Taipei Section

Certificate (Click to expand)

Chi-Mao Fan(student)、Tsung-Jung Liu(adviser)

The pdf file of thesis will publish after 3 years (2025)

Abstract : Image Restoration is a compute vision task which restoring from the degraded images to clean images. With the rapid development of both hardware and software equipment, convolution neural network (CNN) which needs higher equipment requirements gradually replaces the traditional algorithm-based restoration methods, and shine in each domain of compute vision and achieve the impressive performance. However, with the rapid development of hardware equipment, the trend of convolution neural network gradually becomes a clammy situation which stacking huge complex network structures to achieve excellent performances, while ignoring the efficiency of model. In this paper, we based on light hierarchical network architecture: U-Net, and improve from Residual Dense Block (RDB) which is good at image restoration tasks but memory-consuming to an efficient block called Selective Residual Block (SRB). We also improve the hierarchical network structure U-Net by adding the gatepost feature paths which enrich more spatial feature information comparing with the traditional U-Net and have the synergy with SRB. Besides this, we also proposed a loss function which is based on two important metrics in image restoration: peak signal-to-noise (PSNR) and structural similarity index to optimize our model. Finally, proposed network could handle the several restoration tasks such as denoising, deblurring, deraining, dehazing and low-light image enhancement. Furthermore, the performances are good in terms of quantitative metrics and visual quality.

Network Architecture

Overall Framework of SRMNet
Selective Residual Block (SRB)	Resizing Block (Pixel Shuffle)

Installation

The model is built in PyTorch 1.8.0 and tested on Windows10 environment
(Python: 3.8, CUDA: 10.2, cudnn: 7.6).

For installing, follow these intructions

conda create -n pytorch python=3.8
conda activate pytorch
conda install pytorch=1.8 torchvision cudatoolkit=10.2 -c pytorch
conda install -c conda-forge tensorboardx
pip install -r requirements.txt

Quick run

You can directly test your own images on my spcae of HuggingFace.
To test the pre-trained models of SRMNet on your own images, make sure you have downloaded the pre-trained models and place to ./pretrained_model/.
After downloading the models, see quick_demo.py and run

python quick_demo.py --input_dir sample_images --result_dir sample_results --weights_root pretrained_model --gpu True --task [restoration tasks]

All pre-trained models can be downloaded at pretrained_model/README.md or here.

Prepare datasets

The preparation of dataset in more detail, see datasets/README.md.
In our experiments, we crop both training & testing data with the size of 256x256 by the code: generate_patches.py.
More details aboult different restoration tasks can be found in section 4.3 of my thesis.

Train

To train the restoration models of SRMNet. You should check the following components are correct:

training.yaml (Take low-light image enhancement for example)

 # Training configuration
 GPU: [0,1,2,3]

 VERBOSE: False

 MODEL:
   MODE: 'Enhancement'

 # Optimization arguments.
 OPTIM:
   BATCH: 2
   EPOCHS: 200
   # EPOCH_DECAY: [10]
   LR_INITIAL: 2e-4
   LR_MIN: 1e-5
   # BETA1: 0.9

 TRAINING:
   VAL_AFTER_EVERY: 1
   RESUME: False
   TRAIN_PS: 256
   VAL_PS: 256
   TRAIN_DIR: './datasets/train/LOL/train'       # path to training data
   VAL_DIR: './datasets/train/LOL/test' # path to validation data
   SAVE_DIR: './checkpoints'           # path to save models and images

If the data path and above configuration are all correctly setting, just simply run:

python train.py

Test

To test (evaluate) the quantitative scores of different tasks, see evaluation_code.

Results

We achieved competitive performance on image denoising (gaussian & real-world), image deblurring, image deraining (rainstreak & raindrop), image dehazing and low-light image enhancement. Detailed results can be found in section 4.3 of my thesis.

Denoise (gaussian noise) (click me)

Denoise (real-world noise)

Deblur (motion blur)

Derain (rain streak)

Derain (raindrop)

Dehaze (dense haze)

Enhancement (low-light)

Visual Performance

We achieved competitive performance on image denoising (gaussian & real-world), image deblurring, image deraining (rainstreak & raindrop), image dehazing, low-light image enhancement and image retouching. The following figures are some visual performances for different tasks. More detailed results can also be found in section 4.3 of my thesis or directly see in the interactive website.

Denoise (gaussian noise) (click me)

Denoise (real-world noise)

Deblur (motion blur)

Derain (rain streak)

Derain (raindrop)

Dehaze (dense haze)

Enhancement (low-light)

Retouching

Contact

If you have any question, feel free to contact qaz5517359@gmail.com.