CAT-main.zip

# Cross Aggregation Transformer for Image Restoration Zheng Chen, [Yulun Zhang](http://yulunzhang.com/), [Jinjin Gu](https://www.jasongt.com/), [Yongbing Zhang](https://shinning0821.github.io/index.html), [Linghe Kong](https://www.cs.sjtu.edu.cn/~linghe.kong/), and [Xin Yuan](https://xygroup6.github.io/xygroup/), "Cross Aggregation Transformer for Image Restoration", NeurIPS, 2022 (Spotlight) [[paper](https://openreview.net/pdf?id=wQ2QNNP8GtM)] [[arXiv]](https://arxiv.org/abs/2211.13654) [[supplementary material](https://openreview.net/attachment?id=wQ2QNNP8GtM&name=supplementary_material)] [[visual results](https://drive.google.com/drive/folders/1SIQ342yyrlHTCxINf9wYNchOa5eOw_7s?usp=sharing)] [[pretrained models](https://drive.google.com/drive/folders/1ebwl3ahPFczEswRNNIKYbO2PwZt0ZbqU?usp=sharing)] --- > **Abstract:** *Recently, Transformer architecture has been introduced into image restoration to replace convolution neural network (CNN) with surprising results. Considering the high computational complexity of Transformer with global attention, some methods use the local square window to limit the scope of self-attention. However, these methods lack direct interaction among different windows, which limits the establishment of long-range dependencies. To address the above issue, we propose a new image restoration model, Cross Aggregation Transformer (CAT). The core of our CAT is the Rectangle-Window Self-Attention (Rwin-SA), which utilizes horizontal and vertical rectangle window attention in different heads parallelly to expand the attention area and aggregate the features cross different windows. We also introduce the Axial-Shift operation for different window interactions. Furthermore, we propose the Locality Complementary Module to complement the self-attention mechanism, which incorporates the inductive bias of CNN (e.g., translation invariance and locality) into Transformer, enabling global-local coupling. Extensive experiments demonstrate that our CAT outperforms recent state-of-the-art methods on several image restoration applications.* > > <p align="center"> > <img width="800" src="figs/git.png"> > </p> --- | SR (x4) | HQ | LQ | [SwinIR](https://github.com/JingyunLiang/SwinIR) | CAT (ours) | | :--------------------------------------------------: | :-------------------------------------------: | :------------------------------------------------: | :-----------------------------------------------: | :--------------------------------------------: | | <img src="figs/img_024_x4.png" height=80 width=110/> | <img src="figs/img_024_HR_x4.png" height=80/> | <img src="figs/img_024_Bicubic_x4.png" height=80/> | <img src="figs/img_024_SwinIR_x4.png" height=80/> | <img src="figs/img_024_CAT_x4.png" height=80/> | | <img src="figs/img_074_x4.png" height=80 width=110/> | <img src="figs/img_074_HR_x4.png" height=80/> | <img src="figs/img_074_Bicubic_x4.png" height=80/> | <img src="figs/img_074_SwinIR_x4.png" height=80/> | <img src="figs/img_074_CAT_x4.png" height=80/> | ## Dependencies - Python 3.8 - PyTorch 1.8.0 - NVIDIA GPU + [CUDA](https://developer.nvidia.com/cuda-downloads) ```bash # Clone the github repo and go to the default directory 'CAT'. git clone https://github.com/zhengchen1999/CAT.git pip install -r requirements.txt python setup.py develop ``` ## TODO * [x] Image SR * [x] JPEG Compression Artifact Reduction * [x] Image Denoising * [ ] Other tasks ## Contents 1. [Datasets](#Datasets) 1. [Models](#Models) 1. [Training](#Training) 1. [Testing](#Testing) 1. [Results](#Results) 1. [Citation](#Citation) 1. [Acknowledgements](#Acknowledgements) --- ## Datasets Used training and testing sets can be downloaded as follows: | Task | Training Set | Testing Set | Visual Results | | :-------------------------------------------- | :----------------------------------------------------------: | :----------------------------------------------------------: | :----------------------------------------------------------: | | image SR | [DIV2K](https://data.vision.ee.ethz.ch/cvl/DIV2K/) (800 training images) + [Flickr2K](https://cv.snu.ac.kr/research/EDSR/Flickr2K.tar) (2650 images) [complete training dataset [DF2K](https://drive.google.com/file/d/1TubDkirxl4qAWelfOnpwaSKoj3KLAIG4/view?usp=share_link)] | Set5 + Set14 + BSD100 + Urban100 + Manga109 [complete testing dataset [download](https://drive.google.com/file/d/1yMbItvFKVaCT93yPWmlP3883XtJ-wSee/view?usp=sharing)] | [here](https://drive.google.com/drive/folders/122LBzNSuc-YwzyTzA2VL9mXSJPWBxICZ?usp=sharing) | | grayscale JPEG compression artifact reduction | [DIV2K](https://data.vision.ee.ethz.ch/cvl/DIV2K/) (800 training images) + [Flickr2K](https://cv.snu.ac.kr/research/EDSR/Flickr2K.tar) (2650 images) + [WED](http://ivc.uwaterloo.ca/database/WaterlooExploration/exploration_database_and_code.rar)(4744 images) + [BSD500](http://www.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/BSR/BSR_bsds500.tgz) (400 training&testing images) [complete training dataset [DFWB](https://drive.google.com/file/d/1IASyJRsX9CKBE0i5iSJMelIr_a6U5Qcd/view?usp=share_link)] | Classic5 +LIVE + Urban100 [complete testing dataset [download](https://drive.google.com/file/d/17hwSblurN93ndKFRFQQdoRgB-6pFmGtd/view?usp=sharing)] | [here](https://drive.google.com/drive/folders/1xwBMPRUIAnpjAynEr9GI8Wenl6D-J8i3?usp=sharing) | | real image denoising | [SIDD](https://www.eecs.yorku.ca/~kamel/sidd/) (320 training images) [complete training dataset [SIDD](https://drive.google.com/drive/folders/1L_8ig1P71ikzf8PHGs60V6dZ2xoCixaC?usp=share_link)] | [SIDD](https://drive.google.com/file/d/11vfqV-lqousZTuAit1Qkqghiv_taY0KZ/view?usp=sharing) + [DND](https://drive.google.com/file/d/1CYCDhaVxYYcXhSfEVDUwkvJDtGxeQ10G/view?usp=sharing) [complete testing dataset [download](https://drive.google.com/file/d/1Vuu0uhm_-PAG-5UPI0bPIaEjSfrSvsTO/view?usp=share_link)] | [here](https://drive.google.com/drive/folders/14chIIFh6uG4M-aOyJcu6mYjDIpm4zE5t?usp=sharing) | Here the visual results are generated under SR (x4), JPEG compression artifact reduction (q10), and real image denoising. Download training and testing datasets and put them into the corresponding folders of `datasets/` and `restormer/datasets`. See [datasets](datasets/README.md) for the detail of directory structure. ## Models | Task | Method | Params (M) | FLOPs (G) | Dataset | PSNR (dB) | SSIM | Model Zoo | Visual Results | | :-----: | :------ | :--------: | :-------: | :------: | :-------: | :----: | :----------------------------------------------------------: | :----------------------------------------------------------: | | SR | CAT-R | 16.60 | 292.7 | Urban100 | 27.45 | 0.8254 | [Google Drive](https://drive.google.com/drive/folders/1xAFaLnUyloWxK-aXwGzF3rYK0kVh87G5?usp=sharing) | [Google Drive](https://drive.google.com/file/d/1_8ybr6WWvX5YXc7aluq31ZE12mpQ5upl/view?usp=share_link) | | SR | CAT-A | 16.60 | 360.7 | Urban100 | 27.89 | 0.8339 | [Google Drive](https://drive.google.com/drive/folders/1Xm4xQXI74gZcPwgmQHw1qbgdEA0kVCSP?usp=sharing) | [Google Drive](https://drive.google.com/file/d/1vO04maPnzPMJ7y8B8u3-fZsbGQYR6-Hn/view?usp=share_link) | | SR | CAT-R-2 | 11.93 | 216.3 | Urban100 | 27.59 | 0.8285 | [Google Drive](https://drive.google.com/drive/folders/175wdTqjpURS7TSRIj3DODVI_ppktqdJu?usp=sharing) | [Google Drive](https://drive.google.com/file/d/1jSZMdlOrHaiaCLEqP_xAy6hPKhM-McX4/view?usp=share_link)