Skip to content
/ MCLIR Public

[AAAI 2024] Learning from History: Task-agnostic Model Contrastive Learning for Image Restoration

Notifications You must be signed in to change notification settings

Aitical/MCLIR

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

7 Commits
 
 

Repository files navigation

Learning from History: Task-agnostic Model Contrastive Learning for Image Restoration

Gang Wu (吴刚), Junjun Jiang (江俊君), Kui Jiang (江奎), and Xianming Liu (刘贤明)

AIIA Lab, Harbin Institute of Technology, Harbin 150001, China.

Paper arXiv Hits

Contrastive learning has emerged as a prevailing paradigm for high-level vision tasks, which, by introducing properly negative samples, has also been exploited for low-level vision tasks to achieve a compact optimization space to account for their ill-posed nature. However, existing methods rely on manually predefined and task-oriented negatives, which often exhibit pronounced task-specific biases. To address this challenge, our paper introduces an innovative method termed 'learning from history', which dynamically generates negative samples from the target model itself. Our approach, named Model Contrastive paradigm for Image Restoration (MCIR), rejuvenates latency models as negative models, making it compatible with diverse image restoration tasks. We propose the Self-Prior guided Negative loss (SPN) to enable it. This approach significantly enhances existing models when retrained with the proposed model contrastive paradigm. The results show significant improvements in image restoration across various tasks and architectures. For example, models retrained with SPN outperform the original FFANet and DehazeFormer by 3.41 dB and 0.57 dB on the RESIDE indoor dataset for image dehazing. Similarly, they achieve notable improvements of 0.47 dB on SPA-Data over IDT for image deraining and 0.12 dB on Manga109 for a 4x scale super-resolution over lightweight SwinIR, respectively.

Model Contrastive Paradigm for Image Restoration

frameworkv2_new.jpeg

📚Compared to Previous Methods

Methods Task & Dataset PSNR SSIM
FFANet (Qin et al. 2020b) 36.39 0.9886
+CR (Wu et al. 2021) Image Dehazing 36.74 0.9906
+CCR (Zheng et al. 2023) (SOTS-indoor) 39.24 0.9937
+SPN (Ours) 39.80 0.9947
EDSR (Lim et al. 2017) 26.04 0.7849
+PCL (Wu, Jiang, and Liu 2023) SISR 26.07 0.7863
+SPN (Ours) (Urban100) 26.12 0.7878

🔥Easy to Follow

There is a simple implementtation of our Model Contrastive Paradigm and Self-Prior Guided Negative Loss.

-def train_iter(model, lq_input, hq_output, current_iter)
+def train_iter(model, negative_model, lq_input, hq_output, current_iter, lambda, update_step):
    optimizer.zero_grad()
    output = target_model(lq_input)
    L_rec = l1_loss(output, hq_gt)

+   ## Add Negative Sample
+   neg_sample = negative_model(lq_input)
+   ## Add Negative Loss
+   L_neg = perceptual_vgg_loss(output, neg_sample)
+   Loss = L_rec + lambda * L_neg
    Loss.backward()
    optimizer.step()
+   if current_iter % update_step == 0:
+    update_model_ema(negative_model, target_model)

With just a few modifications to your own training scripts, you can easily integrate our approach. Enjoy it!

Results

improvementv2.jpeg

Image Super-Resolution

Method Architecture Scale Avg. Set14 B100 Urban100 Manga109
PSNR/SSIM PSNR/SSIM PSNR/SSIM PSNR/SSIM PSNR/SSIM
EDSR-light x2 32.06/0.9303 33.57/0.9175 32.16/0.8994 31.98/0.9272 30.54/0.9769
+SPN (Ours) CNN 32.19/0.9313 33.67/0.9182 32.21/0.9001 32.23/0.9297 30.64/0.9772
EDSR-light x4 28.14/0.8021 28.58/0.7813 27.57/0.7357 26.04/0.7849 30.35/0.9067
+SPN (Ours) 28.21/0.8040 28.63/0.7829 27.59/0.7369 26.12/0.7878 30.51/0.9085
SwinIR-light x4 28.46/0.8099 28.77/0.7858 27.69/0.7406 26.47/0.7980 30.92/0.9151
+SPN (Ours) Transformer 28.55/0.8114 28.85/0.7874 27.72/0.7414 26.57/0.8010 31.04/0.9158
SwinIR x4 28.88/0.8190 28.94/0.7914 27.83/0.7459 27.07/0.8164 31.67/0.9226
+SPN (Ours) 28.93/0.8198 29.01/0.7923 27.85/0.7465 27.14/0.8176 31.75/0.9229

SR_vision.jpeg
SR_vision_More.jpeg

Image Dehazing

FFANet_Vision.jpeg

Methods SOTS-indoor SOTS-mix
PSNR SSIM PSNR SSIM
(ICCV'19) GridDehazeNet 32.16 0.984 25.86 0.944
(CVPR'20) MSBDN 33.67 0.985 28.56 0.966
(ECCV'20) PFDN 32.68 0.976 28.15 0.962
(AAAI'20) FFANet 36.39 0.989 29.96 0.973
(Ours) FFANet+SPN 39.80 0.995 30.65 0.976
(TIP'23) DehazeFormer-T 35.15 0.989 30.36 0.973
(Ours) DehazeFormer-T+SPN 35.51 0.990 30.44 0.974
(TIP'23)DehazeFormer-S 36.82 0.992 30.62 0.976
(Ours) DehazeFormer-S+SPN 37.24 0.993 30.77 0.978
(TIP'23) DehazeFormer-B 37.84 0.994 31.45 0.980
(Ours) DehazeFormer-B+SPN 38.41 0.994 31.57 0.981

Image Deraining

rain_vision.jpeg A divergence map delineates the differences, highlighting the improvement achieved by ours, particularly in degraded regions.

Method Avg. Rain100L Rain100H DID DDN SPA
undefined PSNR/SSIM PSNR/SSIM PSNR/SSIM PSNR/SSIM PSNR/SSIM PSNR/SSIM
(CVPR21) MPRNet 36.17/0.9543 39.47/0.9825 30.67/0.9110 33.99/0.9590 33.10/0.9347 43.64/0.9844
(AAAI'21) DualGCN 36.69/0.9604 40.73/0.9886 31.15/0.9125 34.37/0.9620 33.01/0.9489 44.18/0.9902
(ICCV'21) SPDNet 36.54/0.9594 40.50/0.9875 31.28/0.9207 34.57/0.9560 33.15/0.9457 43.20/0.9871
(ICCV'21) SwinIR 36.91/0.9507 40.61/0.9871 31.76/0.9151 34.07/0.9313 33.16/0.9312 44.97/0.9890
(CVPR'22) Uformer-S 36.95/0.9505 40.20/0.9860 30.80/0.9105 34.46/0.9333 33.14/0.9312 46.13/0.9913
(CVPR'22) Restormer 37.49/0.9530 40.58/0.9872 31.39/0.9164 35.20/0.9363 34.04/0.9340 46.25/0.9911
(TPAMI'23) IDT 37.77/0.9593 40.74/0.9884 32.10/0.9343 34.85/0.9401 33.80/0.9407 47.34/0.9929
(Ours) IDT+SPN 38.03/0.9610 41.12/0.9893 32.17/0.9352 34.94/0.9424 33.90/0.9442 48.04/0.9938

Image Deblurring

Method MIMO-UNet HINet MAXIM Restormer UFormer NAFNet NAFNet+SPN (Ours)
PSNR 32.68 32.71 32.86 32.92 32.97 32.87 32.93
SSIM 0.959 0.959 0.961 0.961 0.967 0.9606 0.9619

Retrained Models

EDSR baseline SwinIR-light SwinIR-Large FFANet
Download Download Download Download

Quick Evaluation Guide

For quickly evaluating, download the retrained models enhanced by our Model Contrastive Learning. The test scripts for each model are available in their respective repositories: BasicSR, FFANet, DehazeFormer, IDT, and NAFNet. Our gratitude goes out to the authors for their nice sharing of these projects.

Reference

  • Xu Qin, Zhilin Wang, Yuanchao Bai, Xiaodong Xie, Huizhu Jia: FFA-Net: Feature Fusion Attention Network for Single Image Dehazing. AAAI 2020: 11908-11915
  • Bee Lim, Sanghyun Son, Heewon Kim, Seungjun Nah, Kyoung Mu Lee: Enhanced Deep Residual Networks for Single Image Super-Resolution. CVPR Workshops 2017: 1132-1140
  • G. Wu, J. Jiang and X. Liu, "A Practical Contrastive Learning Framework for Single-Image Super-Resolution," in IEEE Transactions on Neural Networks and Learning Systems, doi: 10.1109/TNNLS.2023.3290038
  • Haiyan Wu, Yanyun Qu, Shaohui Lin, Jian Zhou, Ruizhi Qiao, Zhizhong Zhang, Yuan Xie, Lizhuang Ma: Contrastive Learning for Compact Single Image Dehazing. CVPR 2021: 10551-10560
  • Yu Zheng, Jiahui Zhan, Shengfeng He, Junyu Dong, Yong Du: Curricular Contrastive Regularization for Physics-aware Single Image Dehazing. CVPR 2023

Citation

If you find this project useful, please consider citing:

@inproceedings{MCLIR,
  author       = {Gang Wu and
                  Junjun Jiang and
                  Kui Jiang and
                  Xianming Liu},
  title        = {Learning from History: Task-agnostic Model Contrastive Learning for
                  Image Restoration},
  booktitle    = {AAAI},
  year         = {2024}
}

About

[AAAI 2024] Learning from History: Task-agnostic Model Contrastive Learning for Image Restoration

Topics

Resources