iiTransformer: A Unified Approach to Exploiting Local and Non-Local Information for Image Restoration
(PDF,
Supplemental,
Poster,
Video)
Soo Min Kang, Youngchan Song, Hanul Shin, and Tammy Lee
Samsung Research
This is an official repository for iiTransformer (BMVC2022) (Spotlight).
iiTransformer is a general purpose backbone architecture for image restoration, where the goal is to recover a high quality (HQ) image from its degraded low quality (LQ) input. Some examples of image restoration include image denoising, compression artifact removal, and single image super-resolution.
Our approach exploits local and non-local information by utilizing information within the local vicinity of a degraded pixel as well as the data recurrence tendency in natural images. To this end, we use Transformers to capture long-range dependencies at the pixel- or patch-level by switching between intra and inter MSA in the Transformer layer with a reshape operation. We also provide a solution to support inferencing images with arbitrary resolutions without image rescaling or applying the sliding window approach.
We demonstrate the effectiveness of iiTransformer on 3 restoration tasks: image denoising, compression artifact removal (CAR), and single image super-resolution (SISR).
During training, LQ images for the image denoising task are generated on-the-fly, while LQ images of CAR and SISR are prepared before-hand. To ensure consistency across experiments during inferece time, we prepared LQ images of image denoising before-hand. We provide code to prepare LQ images of each task in the data_preparation
folder; refer to the table below for a summary of the file corresponding to each task. For each code, specify the path of HQ/GT and target LQ path in the SRC_PATH
and DST_PATH
variables, respectively.
Task | Script Filename | Language | |
---|---|---|---|
1 | Image Denoising | generate_lq_nr.py | PyTorch |
2 | Compression Artifact Removal (CAR) | generate_lq_car.m | MATLAB |
3 | Single Image Super-resolution (SISR) | generate_lr.m | MATLAB |
To train iiTransformer,
- prepare configuration file in the
options
folder (for reference, see sample json files for each task in theoptions
folder), then - run the following training code:
python3 -m torch.distributed.launch --nproc_per_node=8 --master_port=1234 main_train.py --opt options/train_ii_denoise_color_25.json --dist True
Successful progress in training should create a path {root}/{task}
according to the config file (e.g., denoising/ii_denoise_color_25
) then further create 3 additional folders: images, models, and options. The images
and models
directories will contain images from the validation set and model weights, respectively, at every checkpoint_save
; the options
directory will contain the config file that was used to train the model.
To test the effectiveness of iiTransformer,
- specify the path to the trained weight as
model_path
. - Specify the path to the config file as
json_file
. - Specify the path containing the LQ and HQ images as
lr_path
andhr_path
, respectively. - Specify the path to save the restored images,
save_path
. - Specify the task as one of the following:
classical_sr
,color_dn
, orjpeg_car
and its corresponding degradation level. That is,- if the task is
color_dn
, use thenoise
option to specify the noise level; - if the task is
jpeg_car
, use thejpeg
option to specify the compression level; - if the task is
classical_sr
, then use thescale
option to specify the scale.
- if the task is
- Run the following inference code:
python3 main_test.py --task {task} --{degradation_type} {degradation_level} --model_path ${model_path} --folder_lq ${lr_path} --folder_gt ${hr_path} --save_path ${save_path} --tile 128 --tile_overlap 64 --opt ${json_file}
See below for an example of full inference code restoring noisy images with AWGN of sigma=25:
model_path='denoising/ii_denoise_color_25/models/460000_E.pth'
json_file='denoising/ii_denoise_color_25/options/train_ii_denoise_color_25_220506_121143.json'
lr_path='DB/Kodak24/DN/s25_torch'
hr_path='DB/Kodak24/GT'
save_path='denoising/ii_denoise_color_25/results/Kodak24'
python3 main_test.py --task color_dn --noise 25 --model_path ${model_path} --folder_lq ${lr_path} --folder_gt ${hr_path} --save_path ${save_path} --tile 128 --tile_overlap 64 --opt ${json_file}
Please cite our paper if you use this code or our model:
@inproceedings{Kang_2022_BMVC,
author = {Soo Min Kang and Youngchan Song and Hanul Shin and Tammy Lee},
title = {iiTransformer: A Unified Approach to Exploiting Local and Non-local Information for Image Restoration},
booktitle = {33rd British Machine Vision Conference 2022, {BMVC} 2022, London, UK, November 21-24, 2022},
publisher = {{BMVA} Press},
year = {2022},
url = {https://bmvc2022.mpi-inf.mpg.de/0377.pdf}
}
This repository was built on the Image Restoration Toolbox (PyTorch).