This is the official PyTorch codes for the paper
Correlation Matching Transformation Transformers for UHD Image Restoration
Cong Wang, Jinshan Pan, Wei Wang, Gang Fu, Siyuan Liang, Mengzhu Wang, Xiao-Ming Wu, Jun Liu
This paper proposes UHDformer, a general Transformer for Ultra-High-Definition (UHD) image restoration.
UHDformer contains two learning spaces: (a) learning in high-resolution space and (b) learning in low-resolution space.
The former learns multi-level high-resolution features and fuses low-high features and reconstructs the residual images, while the latter
explores more representative features learning from the highresolution ones to facilitate better restoration.
To better improve feature representation in low-resolution space, we propose to build feature transformation from the high-resolution space to the low-resolution one.
To that end, we propose two new modules: Dual-path Correlation Matching Transformation module (DualCMT) and Adaptive Channel Modulator (ACM).
The DualCMT selects top C/r (r is greater or equal to 1 which controls the squeezing level) correlation channels
from the max-pooling/mean-pooling high-resolution features to replace low-resolution ones in Transformers, which can effectively
squeeze useless content to improve the feature representation in low-resolution space to facilitate better recovery.
The ACM is exploited to adaptively modulate multi-level high-resolution features, enabling to provide more useful features to low-resolution space for better learning.
Experimental results show that our UHDformer reduces about ninetyseven percent model sizes compared with most state-of-theart methods while significantly improving performance under different training sets on 3 UHD image restoration tasks, including
low-light image enhancement, image dehazing, and image deblurring.
- Ubuntu >= 18.04
- CUDA >= 11.0
- Other required packages in
requirements.txt
# git clone this repository
git clone https://github.com/supersupercong/UHDformer.git
cd UHDformer
# create new anaconda env
conda create -n uhdformer python=3.8
source activate uhdformer
# install python dependencies
pip3 install -r requirements.txt
python setup.py develop
UHD-LL, UHD-Haze [coming soon], UHD-Blur [coming soon]
UHD-LL&UHD-Haze&UHD-Blur&LOL-SOTS
bash train.sh
bash test.sh
@inproceedings{wang2024uhdformer,
author={Cong Wang and Jinshan Pan and Wei Wang and Gang Fu and Siyuan Liang and Mengzhu Wang and Xiao-Ming Wu and Jun Liu},
title={Correlation Matching Transformation Transformers for UHD Image Restoration},
year={2024},
Journal = {Proceedings of the AAAI Conference on Artificial Intelligence (AAAI)},
}
This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License.
Any questions can contact: Cong Wang [supercong94@gmail.com]
This project is based on FeMaSR.