[NCHU thesis]
Image Restoration by Selective Residual Block on Improved Hierarchical Encoder-Decoder Networks
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.
Overall Framework of SRMNet |
|
Selective Residual Block (SRB) |
Resizing Block (Pixel Shuffle) |
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
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.
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.
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
To test (evaluate) the quantitative scores of different tasks, see evaluation_code
.
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.
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.
If you have any question, feel free to contact [email protected].