Awesome
Learning Multiple Adverse Weather Removal via Two-stage Knowledge Learning and Multi-contrastive Regularization: Toward a Unified Model
[CVPR2022] Official Pytorch based implementation.
<hr />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
<table> <tr> <td colspan="2" align="center"> <img src = "https://github.com/fingerk28/Two-stage-Knowledge-For-Multiple-Adverse-Weather-Removal/blob/main/images/architecture.png"> </td> </tr> <tr> <td colspan="2" align="center"><p><b>Overall Architecture</b></p></td> </tr> <tr> <td align="center"> <img src = "https://github.com/fingerk28/Two-stage-Knowledge-For-Multiple-Adverse-Weather-Removal/blob/main/images/ckt.png" width="400"> </td> <td align="center"> <img src = "https://github.com/fingerk28/Two-stage-Knowledge-For-Multiple-Adverse-Weather-Removal/blob/main/images/mcr.png" width="400"> </td> </tr> <tr> <td align="center"><p><b>Collaborative Knowledge Trasfer</b></p></td> <td align="center"><p><b>Multi-contrastive Regularization</b></p></td> </tr> </table>Quantitative Result
<table> <tr> <td align="center"> <img src = "https://github.com/fingerk28/Two-stage-Knowledge-For-Multiple-Adverse-Weather-Removal/blob/main/images/quantitative_result.png" width="400"> </td> </tr> <tr> <td align="center"><p><b>Setting1</b></p></td> </tr> </table>Qualitative Result
<table> <tr> <td align="center"> <img src = "https://github.com/fingerk28/Two-stage-Knowledge-For-Multiple-Adverse-Weather-Removal/blob/main/images/qualitative_result.png"> </td> </tr> <tr> <td align="center"><p><b>Setting1</b></p></td> </tr> </table>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
andDatasetForValid
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}
}