Multi-Stage Progressive Image Restoration

Overview

PWC PWC PWC PWC PWC PWC PWC PWC PWC PWC PWC

Multi-Stage Progressive Image Restoration

Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, Ming-Hsuan Yang, and Ling Shao

Paper: https://arxiv.org/abs/2102.02808

Supplementary: pdf

Abstract: Image restoration tasks demand a complex balance between spatial details and high-level contextualized information while recovering images. In this paper, we propose a novel synergistic design that can optimally balance these competing goals. Our main proposal is a multi-stage architecture, that progressively learns restoration functions for the degraded inputs, thereby breaking down the overall recovery process into more manageable steps. Specifically, our model first learns the contextualized features using encoder-decoder architectures and later combines them with a high-resolution branch that retains local information. At each stage, we introduce a novel per-pixel adaptive design that leverages in-situ supervised attention to reweight the local features. A key ingredient in such a multi-stage architecture is the information exchange between different stages. To this end, we propose a two-faceted approach where the information is not only exchanged sequentially from early to late stages, but lateral connections between feature processing blocks also exist to avoid any loss of information. The resulting tightly interlinked multi-stage architecture, named as MPRNet, delivers strong performance gains on ten datasets across a range of tasks including image deraining, deblurring, and denoising. For example, on the Rain100L, GoPro and DND datasets, we obtain PSNR gains of 4 dB, 0.81 dB and 0.21 dB, respectively, compared to the state-of-the-art.

Network Architecture

Overall Framework of MPRNet

Supervised Attention Module (SAM)

Installation

The model is built in PyTorch 1.1.0 and tested on Ubuntu 16.04 environment (Python3.7, CUDA9.0, cuDNN7.5).

For installing, follow these intructions

conda create -n pytorch1 python=3.7
conda activate pytorch1
conda install pytorch=1.1 torchvision=0.3 cudatoolkit=9.0 -c pytorch
pip install matplotlib scikit-image opencv-python yacs joblib natsort h5py tqdm

Evaluation

Testing codes for deblurring, deraining and denoising are provided in their respective directories.

Results

Experiments are performed for different image processing tasks including, image deblurring, image deraining and image denoising.

Image Deblurring

Deblurring on Synthetic Datasets.

Deblurring on Real Dataset.

Image Deraining

Image Denoising

Citation

If you use MPRNet, please consider citing:

@article{Zamir2021MPRNet,
    title={Multi-Stage Progressive Image Restoration},
    author={Syed Waqas Zamir and Aditya Arora and Salman Khan and Munawar Hayat
            and Fahad Shahbaz Khan and Ming-Hsuan Yang and Ling Shao},
    journal={arXiv:2102.02808},
    year={2020}
}

Contact

Should you have any question, please contact [email protected]

Comments
  • RuntimeError: CUDA out of memory

    RuntimeError: CUDA out of memory

    While trying/testing a pretrained model with the cmd:

    python demo.py --task Deblurring --input_dir path_to_images --result_dir save_images_here
    

    I am getting some runtime error. So basically, the terminal shows this

    RuntimeError: CUDA out of memory. Tried to allocate 2.15 GiB (GPU 0; 1.96 GiB total capacity; 420.11 MiB already allocated; 400.00 MiB free; 233.89 MiB cached)
    
    

    I am guessing this is most likely because my GPU's memory is occupied and is not sufficient for some allocation that the program tries during runtime. Since you can see that the allocation size if already greater than the total capacity, so could you suggest or recommend some methods or change in the code for reducing this allocation size. Thanks

    opened by uV3301 14
  • How much CUDA memory is required according to your default Settings?

    How much CUDA memory is required according to your default Settings?

    when i'm running your codes,it's out of memory on our RTX3090 with 24G memory. we can only run it with batch 16 and patch 128 on tow GPU(15G used per GPU) how should i balance the patch size and batch size. thank you!

    opened by zhangheng-h 13
  • some issue about the data loader~

    some issue about the data loader~

    Thank you very much for sharing, this is a very interesting paper. I face some issues when I try to train the denoise and Derain model. when I run the Derain the train.py and download the dataset.i face the "RuntimeError: stack expects each tensor to be equal size, but got [3, 229, 256] at entry 0 and [3, 256, 171] at entry 1". should I to pretrain the data in the same size before I run the train.py of the Derain the model? could you give me a hand for this problem? thanks very much

    opened by leefly072 10
  • setting of pre-train model

    setting of pre-train model

    Hi Thanks for the great work! I am wondering what is the setting of the pre-trained model that you provided for deblurring? Is it n_feat=40, scale_unetfeats=40, scale_orsnetfeats=32 ? Or n_feat=96, scale_unetfeats=48, scale_orsnetfeats=32 ?

    Do you have them both? Are they suitable for PyTroch 1.8?

    Thanks!

    opened by hhcs9527 7
  • AttributeError: 'CosineAnnealingLR' object has no attribute 'get_last_lr'

    AttributeError: 'CosineAnnealingLR' object has no attribute 'get_last_lr'

    when I run deblurring task python train.py , I get the error "AttributeError: 'CosineAnnealingLR' object has no attribute 'get_last_lr'" when the epoch is 3, how can i solve the problem

    opened by kzq666666 7
  • RealBlur data for training

    RealBlur data for training

    Hi, I'd like to know if you have pre-processed the data of RealBlur datasethttp://cg.postech.ac.kr/research/realblur/. Since I directly use those data for training, the loss becomes inf. The learning rate and batch size are as default settings.
    I'd like to know how you select the data for training.Thanks!

    opened by asdf1996 7
  • This problem occurs when training deraining dataset. Why is that?

    This problem occurs when training deraining dataset. Why is that?

    Hello,I'm sorry to bother you. I encountered a problem when training the deraining dataset. My friends didn't encountered this case, so I'd like to ask you to help me see the reason for this error. I only changed the values of BATCH_SIZE and TRAIN_PS in the train. yml file, it occurred during the training of epochs 9, 22, 31 and 60. The error code is as follows:

    Traceback (most recent call last):
      File "train.py", line 109, in <module>
        for i, data in enumerate(tqdm(train_loader), 0):
      File "/home/min1/anaconda3/envs/bwq/lib/python3.7/site-packages/tqdm/std.py", line 1185, in __iter__
        for obj in iterable:
      File "/home/min1/anaconda3/envs/bwq/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 582, in __next__
        return self._process_next_batch(batch)
      File "/home/min1/anaconda3/envs/bwq/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 608, in _process_next_batch
        raise batch.exc_type(batch.exc_msg)
    RuntimeError: Traceback (most recent call last):
      File "/home/min1/anaconda3/envs/bwq/lib/python3.7/site-packages/torch/utils/data/_utils/worker.py", line 99, in _worker_loop
        samples = collate_fn([dataset[i] for i in batch_indices])
      File "/home/min1/anaconda3/envs/bwq/lib/python3.7/site-packages/torch/utils/data/_utils/collate.py", line 68, in default_collate
        return [default_collate(samples) for samples in transposed]
      File "/home/min1/anaconda3/envs/bwq/lib/python3.7/site-packages/torch/utils/data/_utils/collate.py", line 68, in <listcomp>
        return [default_collate(samples) for samples in transposed]
      File "/home/min1/anaconda3/envs/bwq/lib/python3.7/site-packages/torch/utils/data/_utils/collate.py", line 43, in default_collate
        return torch.stack(batch, 0, out=out)
    RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 0. Got 128 and 126 in dimension 3 at /opt/conda/conda-bld/pytorch_1556653114079/work/aten/src/TH/generic/THTensor.cpp:711
    
    opened by lonely6628 6
  • Why is the loss so big in the training?

    Why is the loss so big in the training?

    Deblurring task

    The configuration I changed is as follows: BATCH_SIZE: 16 TRAIN_PS: 64 VAL_PS: 64 why is the training loss so big as the epoch is 140? image

    opened by kzq666666 6
  • model size

    model size

    hello,my name is Long. I downloaded your pretrain model and the model's size is about 80M. When I train a new model,I found the model's size is 241M.what should I do can reduce the size? thank you.

    opened by chenloveheimei 5
  • RuntimeError: The size of tensor a (66) must match the size of tensor b (67) at non-singleton dimension 2

    RuntimeError: The size of tensor a (66) must match the size of tensor b (67) at non-singleton dimension 2

    ===>Testing using weights:  ./pretrained_models/model_denoising.pth
      0%|          | 0/1 [00:00<?, ?it/s]D:\WorkTools\ProgramData\Anaconda3\envs\py36\lib\site-packages\torch\nn\functional.py:3103: UserWarning: The default behavior for interpolate/upsample with float scale_factor changed in 1.6.0 to align with other frameworks/libraries, and now uses scale_factor directly, instead of relying on the computed output size. If you wish to restore the old behavior, please set recompute_scale_factor=True. See the documentation of nn.Upsample for details. 
      warnings.warn("The default behavior for interpolate/upsample with float scale_factor changed "
      0%|          | 0/1 [00:16<?, ?it/s]
    Traceback (most recent call last):
      File "E:/gitrestorys/denoising/MPRNet/Denoising/test_SIDD.py", line 65, in <module>
        restored_patch = model_restoration(noisy_patch)
      File "D:\WorkTools\ProgramData\Anaconda3\envs\py36\lib\site-packages\torch\nn\modules\module.py", line 727, in _call_impl
        result = self.forward(*input, **kwargs)
      File "D:\WorkTools\ProgramData\Anaconda3\envs\py36\lib\site-packages\torch\nn\parallel\data_parallel.py", line 149, in forward
        return self.module(*inputs, **kwargs)
      File "D:\WorkTools\ProgramData\Anaconda3\envs\py36\lib\site-packages\torch\nn\modules\module.py", line 727, in _call_impl
        result = self.forward(*input, **kwargs)
      File "E:\gitrestorys\denoising\MPRNet\Denoising\MPRNet.py", line 300, in forward
        res1_top = self.stage1_decoder(feat1_top)
      File "D:\WorkTools\ProgramData\Anaconda3\envs\py36\lib\site-packages\torch\nn\modules\module.py", line 727, in _call_impl
        result = self.forward(*input, **kwargs)
      File "E:\gitrestorys\denoising\MPRNet\Denoising\MPRNet.py", line 145, in forward
        x = self.up32(dec3, self.skip_attn2(enc2))
      File "D:\WorkTools\ProgramData\Anaconda3\envs\py36\lib\site-packages\torch\nn\modules\module.py", line 727, in _call_impl
        result = self.forward(*input, **kwargs)
      File "E:\gitrestorys\denoising\MPRNet\Denoising\MPRNet.py", line 183, in forward
        x = x + y
    RuntimeError: The size of tensor a (66) must match the size of tensor b (67) at non-singleton dimension 2
    
    Process finished with exit code -1
    

    how to fix it

    opened by Billy4J 5
  • meet Runtime error when training deblur codes

    meet Runtime error when training deblur codes

    Traceback (most recent call last): File "/data/h00026370/Moon/codes/MPRNet/Deblurring/train.py", line 120, in loss_char = np.sum([criterion_char(restored[j],target) for j in range(len(restored))]) File "<array_function internals>", line 6, in sum File "/home/leipeng/anaconda3/envs/mprnet/lib/python3.7/site-packages/numpy/core/fromnumeric.py", line 2260, in sum initial=initial, where=where) File "/home/leipeng/anaconda3/envs/mprnet/lib/python3.7/site-packages/numpy/core/fromnumeric.py", line 86, in _wrapreduction return ufunc.reduce(obj, axis, dtype, out, **passkwargs) File "/home/leipeng/anaconda3/envs/mprnet/lib/python3.7/site-packages/torch/tensor.py", line 458, in array return self.numpy() RuntimeError: Can't call numpy() on Variable that requires grad. Use var.detach().numpy() instead.

    opened by hm-better 4
  • About deraining

    About deraining

    In the paper,you say we compute image quality scores using the Y channel (in YCbCr color space), but in training or testing code, i didn't see it computed on YCBCR color space. Did i missing the code?

    opened by kiriiiia 0
  • Loss raise to abnormal and batchsize

    Loss raise to abnormal and batchsize

    Loss raises to several million after 50 epochs (Before 50 epoch is normal). And why I can only allow batchsize 2 on RTX3090 when training, 2 more will out of memory.

    opened by HongyuZhang140 5
  • About Training, how to deal with this problem?

    About Training, how to deal with this problem?

    Traceback (most recent call last): File "train.py", line 140, in for ii, data_val in enumerate((val_loader), 0): File "/home/mail/2019p9/p930026055/miniconda3/envs/py37/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 530, in next data = self._next_data() File "/home/mail/2019p9/p930026055/miniconda3/envs/py37/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 1204, in _next_data return self._process_data(data) File "/home/mail/2019p9/p930026055/miniconda3/envs/py37/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 1250, in _process_data data.reraise() File "/home/mail/2019p9/p930026055/miniconda3/envs/py37/lib/python3.7/site-packages/torch/_utils.py", line 457, in reraise raise exception IndexError: Caught IndexError in DataLoader worker process 6. Original Traceback (most recent call last): File "/home/mail/2019p9/p930026055/miniconda3/envs/py37/lib/python3.7/site-packages/torch/utils/data/_utils/worker.py", line 287, in _worker_loop data = fetcher.fetch(index) File "/home/mail/2019p9/p930026055/miniconda3/envs/py37/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch data = [self.dataset[idx] for idx in possibly_batched_index] File "/home/mail/2019p9/p930026055/miniconda3/envs/py37/lib/python3.7/site-packages/torch/utils/data/utils/fetch.py", line 49, in data = [self.dataset[idx] for idx in possibly_batched_index] File "/share/mail/2019p9/p930026055/directory_2/MPRNet/Deraining/dataset_RGB.py", line 111, in getitem inp_path = self.inp_filenames[index] IndexError: list index out of range

    opened by jjjjzyyyyyy 2
  • Log Files from Training

    Log Files from Training

    Thank you for your awesome code!

    I am hoping you might open-source the log files you have from training. Maybe the training and validation loss as a function of epoch (and/or batch) with an estimate of the runtime?

    opened by gauenk 0
Owner
Syed Waqas Zamir
Research Scientist
Syed Waqas Zamir
Code for our NeurIPS 2021 paper Mining the Benefits of Two-stage and One-stage HOI Detection

CDN Code for our NeurIPS 2021 paper "Mining the Benefits of Two-stage and One-stage HOI Detection". Contributed by Aixi Zhang*, Yue Liao*, Si Liu, Mia

null 71 Dec 14, 2022
Code for Mining the Benefits of Two-stage and One-stage HOI Detection

Status: Archive (code is provided as-is, no updates expected) PPO-EWMA [Paper] This is code for training agents using PPO-EWMA and PPG-EWMA, introduce

OpenAI 33 Dec 15, 2022
Virtual Dance Reality Stage: a feature that offers you to share a stage with another user virtually

Portrait Segmentation using Tensorflow This script removes the background from an input image. You can read more about segmentation here Setup The scr

null 291 Dec 24, 2022
Code of U2Fusion: a unified unsupervised image fusion network for multiple image fusion tasks, including multi-modal, multi-exposure and multi-focus image fusion.

U2Fusion Code of U2Fusion: a unified unsupervised image fusion network for multiple image fusion tasks, including multi-modal (VIS-IR, medical), multi

Han Xu 129 Dec 11, 2022
Official repository for "Restormer: Efficient Transformer for High-Resolution Image Restoration". SOTA for motion deblurring, image deraining, denoising (Gaussian/real data), and defocus deblurring.

Restormer: Efficient Transformer for High-Resolution Image Restoration Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan,

Syed Waqas Zamir 906 Dec 30, 2022
Code for Referring Image Segmentation via Cross-Modal Progressive Comprehension, CVPR2020.

CMPC-Refseg Code of our CVPR 2020 paper Referring Image Segmentation via Cross-Modal Progressive Comprehension. Shaofei Huang*, Tianrui Hui*, Si Liu,

spyflying 55 Dec 1, 2022
(under submission) Bayesian Integration of a Generative Prior for Image Restoration

BIGPrior: Towards Decoupling Learned Prior Hallucination and Data Fidelity in Image Restoration Authors: Majed El Helou, and Sabine Süsstrunk {Note: p

Majed El Helou 22 Dec 17, 2022
HINet: Half Instance Normalization Network for Image Restoration

HINet: Half Instance Normalization Network for Image Restoration Liangyu Chen, Xin Lu, Jie Zhang, Xiaojie Chu, Chengpeng Chen Paper: https://arxiv.org

null 303 Dec 31, 2022
This is an implementation for the CVPR2020 paper "Learning Invariant Representation for Unsupervised Image Restoration"

Learning Invariant Representation for Unsupervised Image Restoration (CVPR 2020) Introduction This is an implementation for the paper "Learning Invari

GarField 88 Nov 7, 2022
EDPN: Enhanced Deep Pyramid Network for Blurry Image Restoration

EDPN: Enhanced Deep Pyramid Network for Blurry Image Restoration Ruikang Xu, Zeyu Xiao, Jie Huang, Yueyi Zhang, Zhiwei Xiong. EDPN: Enhanced Deep Pyra

null 69 Dec 15, 2022
Real-ESRGAN aims at developing Practical Algorithms for General Image Restoration.

Real-ESRGAN Colab Demo for Real-ESRGAN . Portable Windows executable file. You can find more information here. Real-ESRGAN aims at developing Practica

Xintao 17.2k Jan 2, 2023
Image restoration with neural networks but without learning.

Warning! The optimization may not converge on some GPUs. We've personally experienced issues on Tesla V100 and P40 GPUs. When running the code, make s

Dmitry Ulyanov 7.4k Jan 1, 2023
SwinIR: Image Restoration Using Swin Transformer

SwinIR: Image Restoration Using Swin Transformer This repository is the official PyTorch implementation of SwinIR: Image Restoration Using Shifted Win

Jingyun Liang 2.4k Jan 8, 2023
Image Restoration Toolbox (PyTorch). Training and testing codes for DPIR, USRNet, DnCNN, FFDNet, SRMD, DPSR, BSRGAN, SwinIR

Image Restoration Toolbox (PyTorch). Training and testing codes for DPIR, USRNet, DnCNN, FFDNet, SRMD, DPSR, BSRGAN, SwinIR

Kai Zhang 2k Dec 31, 2022
Dynamic Attentive Graph Learning for Image Restoration, ICCV2021 [PyTorch Code]

Dynamic Attentive Graph Learning for Image Restoration This repository is for GATIR introduced in the following paper: Chong Mou, Jian Zhang, Zhuoyuan

Jian Zhang 84 Dec 9, 2022
An official repository for Paper "Uformer: A General U-Shaped Transformer for Image Restoration".

Uformer: A General U-Shaped Transformer for Image Restoration Zhendong Wang, Xiaodong Cun, Jianmin Bao and Jianzhuang Liu Paper: https://arxiv.org/abs

Zhendong Wang 497 Dec 22, 2022
Image Restoration Using Swin Transformer for VapourSynth

SwinIR SwinIR function for VapourSynth, based on https://github.com/JingyunLiang/SwinIR. Dependencies NumPy PyTorch, preferably with CUDA. Note that t

Holy Wu 11 Jun 19, 2022
Half Instance Normalization Network for Image Restoration

HINet Half Instance Normalization Network for Image Restoration, based on https://github.com/megvii-model/HINet. Dependencies NumPy PyTorch, preferabl

Holy Wu 4 Jun 6, 2022
The 7th edition of NTIRE: New Trends in Image Restoration and Enhancement workshop will be held on June 2022 in conjunction with CVPR 2022.

NTIRE 2022 - Image Inpainting Challenge Important dates 2022.02.01: Release of train data (input and output images) and validation data (only input) 2

Andrés Romero 37 Nov 27, 2022