Learning Multiple Adverse Weather Removal via Two-stage Knowledge Learning and Multi-contrastive Regularization: Toward a Unified Model

[CVPR2022] Official Pytorch based implementation.

[paper]

Abstract: In this paper, an ill-posed problem of multiple adverse weather removal is investigated. Our goal is to train a model with a 'unified' architecture and only one set of pretrained weights that can tackle multiple types of adverse weathers such as haze, snow, and rain simultaneously. To this end, a two-stage knowledge learning mechanism including knowledge collation (KC) and knowledge examination (KE) based on a multi-teacher and student architecture is proposed. At the KC, the student network aims to learn the comprehensive bad weather removal problem from multiple well-trained teacher networks where each of them is specialized in a specific bad weather removal problem. To accomplish this process, a novel collaborative knowledge transfer is proposed. At the KE, the student model is trained without the teacher networks and examined by challenging pixel loss derived by the ground truth. Moreover, to improve the performance of our training framework, a novel loss function called multi-contrastive knowledge regularization (MCR) loss is proposed. Experiments on several datasets show that our student model can achieve promising results on different bad weather removal tasks simultaneously.

Architecture

Overall Architecture
Collaborative Knowledge Trasfer	Multi-contrastive Regularization

Quantitative Result

Setting1

Qualitative Result

Setting1

Usage

Pre-trained Models

• [Setting1] [CSD, Rain1400, ITS, OTS] Download link:

  • CSD Teacher → Google Drive
  • Rain1400 Teacher → Google Drive
  • ITS, OTS Teacher → Google Drive
  • Student → Google Drive

• [Setting2] [Snow100k, Raindrop, Rainfog] Download link:

  • Raindrop Teacher → Google Drive
  • Rainfog Teacher → Google Drive
  • Snow100k Teacher → Google Drive
  • Student → Google Drive

Install

git clone https://github.com/fingerk28/Two-stage-Knowledge-For-Multiple-Adverse-Weather-Removal.git

Training

python train.py --teacher TEACHER_CHECKPOINT_PATH_0 TEACHER_CHECKPOINT_PATH_1 TEACHER_CHECKPOINT_PATH_2 --save-dir RESULTS_WILL_BE_SAVED_HERE

--teacher → input any amout of teacher checkpoint path

---

You need to prepare the meta file (.json) under the ./meta

Class DatasetForTrain and DatasetForValid would take all meta files as the datasources.

The structure should be:

.
├── inference.py
├── meta
│   ├── train
│   │   ├── CSD_meta_train.json
│   │   └── Rain1400_meta_train.json
│   │   └── ...
│   └── valid
│       ├── CSD_meta_valid.json
│       └── Rain1400_meta_valid.json
│       └── ...
├── models
│   ├── ...
├── train.py
└── utils
    ├── ...

---

The structure of the .json file should be:

[
    [
        "path_to_GT_image0",
        "path_to_Input_image0"
    ],
    [
        "path_to_GT_image1",
        "path_to_Input_image1"
    ],
    [
           ...
    ],
    ...

]

Inference

python inference.py --dir_path DIR_OF_TEST_IMAGES --checkpoint CHECKPOINT_PATH --save_dir RESULTS_WILL_BE_SAVED_HERE 

Other Works for Image Restoration

You can also refer to our previous works:

• Desnowing → [JSTASR] (ECCV'20) and [HDCW-Net] (ICCV'21)
• Dehazing → [PMS-Net] (CVPR'19) and [PMHLD] (TIP'20)
• Deraining → [ContouletNet] (BMVC'21)
• Image Relighting → [MB-Net] (NTIRE'21 1st solution) and [S3Net] (NTIRE'21 3 rd solution)

Citation

Please cite this paper in your publications if it is helpful for your tasks.

@inproceedings{Chen2022MultiWeatherRemoval,
  title={Learning Multiple Adverse Weather Removal via Two-stage Knowledge Learning and Multi-contrastive Regularization: Toward a Unified Model},
  author={Chen, Wei-Ting and Huang, Zhi-Kai and Tsai, Cheng-Che and Yang, Hao-Hsiang and Ding, Jian-Jiun and Kuo, Sy-Yen},
  journal={2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2022}
}