Skip to content

Tombs98/CAPTNet

Repository files navigation

Prompt-based Ingredient-Oriented All-in-One Image Restoration

paper


Abstract: Image restoration aims to recover the high-quality images from their degraded observations. Since most existing methods have been dedicated into single degradation removal, they may not yield optimal results on other types of degradations, which do not satisfy the applications in real world scenarios. In this paper, we propose a novel data ingredient-oriented approach that leverages prompt-based learning to enable a single model to efficiently tackle multiple image degradation tasks. Specifically, we utilize a encoder to capture features and introduce prompts with degradation-specific information to guide the decoder in adaptively recovering images affected by various degradations. In order to model the local invariant properties and non-local information for high-quality image restoration, we combined CNNs operations and Transformers. Simultaneously, we made several key designs in the Transformer blocks (multi-head rearranged attention with prompts and simple-gate feed-forward network) to reduce computational requirements and selectively determines what information should be persevered to facilitate efficient recovery of potentially sharp images. Furthermore, we incorporate a feature fusion mechanism further explores the multi-scale information to improve the aggregated features. The resulting tightly interlinked hierarchy architecture, named as CAPTNet, extensive experiments demonstrate that our method performs competitively to the state-of-the-art.

Network Architecture

Overall Framework of CAPTNet

(a) The architecture of non-linear activation free block (NAFBlock). (b) Simplified channel attention (SCA). (c) Simplified prompt-based transformer (SPT) block. (d) Multi-head rearranged attention with prompts (MRAP). (e) Simple gate feed-forward network (SGFN)

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

Install warmup scheduler

cd pytorch-gradual-warmup-lr; python setup.py install; cd ..

Quick Run

To test the pre-trained models Baidu Drive password:Tomb

Google Drive

Training and Evaluation

Training

  • Train the model with default arguments by running
python train.py

Evaluation

  • Run
python test.py 
  • To reproduce PSNR/SSIM scores of the paper , run
python eval.py 

Results

Citations

If our code helps your research or work, please consider citing our paper. The following is a BibTeX reference:

@article{Gao2023PIO,
 author={Gao, Hu and Yang, Jing and Zhang, Ying and Wang, Ning and Yang, Jingfan and Dang, Depeng},
  journal={IEEE Transactions on Circuits and Systems for Video Technology}, 
  title={Prompt-based Ingredient-Oriented All-in-One Image Restoration}, 
  year={2024},
  volume={},
  number={},
  pages={1-1}
}

Contact

Should you have any question, please contact two_bits@163.com

About

Prompt-based Ingredient-Oriented All-in-One Image Restoration

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages