initial code release

README.md

@@ -1,18 +1,51 @@
 # Always Clear Days: Degradation Type and Severity Aware All-In-One Adverse Weather Removal
 An official implement of the paper "Always Clear Days: Degradation Type and Severity Aware All-In-One Adverse Weather Removal"
+<p align="center">
+  <img src="asset/banner.png" width="700"/>
+</p>
 
 [[Paper](https://arxiv.org/abs/2310.18293)]
 
 [Yu-Wei, Chen](https://fordevoted.github.io), [Soo-Chang, Pei](https://scholar.google.com/citations?user=-JiGrnAAAAAJ&hl=zh-TW)
 
+<p align="center">
+  <img src="asset/model.png" width="500"/>
+</p>
 
-> **Abstract:**  All-in-one adverse weather removal is an emerging topic on image restoration, which aims to restore multiple weather degradation in an unified model, and the challenging are twofold. First, discovering and handling the property of multi-domain in target distribution formed by multiple weather conditions. Second, design efficient and effective operations for different degradation types. To address this problem, most prior works focus on the multi-domain caused by weather type. Inspired by inter\&intra-domain adaptation literature, we observed that not only weather type but also weather severity introduce multi-domain within each weather type domain, which is ignored by previous methods, and further limit their performance. To this end, we proposed a degradation type and severity aware model, called UtilityIR, for blind all-in-one bad weather image restoration. To extract weather information from single image, we proposed a novel Marginal Quality Ranking Loss (MQRL) and utilized Contrastive Loss (CL) to guide weather severity and type extraction, and leverage a bag of novel techniques such as Multi-Head Cross Attention (MHCA) and Local-Global Adaptive Instance Normalization (LG-AdaIN) to efficiently restore spatial varying weather degradation. The proposed method can significantly outperform the SOTA methods subjectively and objectively on different weather restoration tasks with a large margin, and enjoy less model parameters. Proposed method even can restore unseen domain combined multiple degradation images, and modulating restoration level.
+
+> **Abstract:**  All-in-one adverse weather removal is an emerging topic on image restoration, which aims to restore multiple weather degradations in an unified model, and the challenge are twofold. First, discover and handle the property of multi-domain in target distribution formed by multiple weather conditions. Second, design efficient and effective operations for different degradations. To resolve this problem, most prior works focus on the multi-domain caused by different weather types. Inspired by inter\&intra-domain adaptation literature, we observe that not only weather type but also weather severity introduce multi-domain within each weather type domain, which is ignored by previous methods, and further limit their performance. To this end, we propose a degradation type and severity aware model, called UtilityIR, for blind all-in-one bad weather image restoration. To extract weather information from single image, we propose a novel Marginal Quality Ranking Loss (MQRL) and utilize Contrastive Loss (CL) to guide weather severity and type extraction, and leverage a bag of novel techniques such as Multi-Head Cross Attention (MHCA) and Local-Global Adaptive Instance Normalization (LG-AdaIN) to efficiently restore spatial varying weather degradation. The proposed method can outperform the state-of-the-art methods subjectively and objectively on different weather removal tasks with a large margin, and enjoy less model parameters. Proposed method even can restore unseen combined multiple degradation images, and modulate restoration level.
 
 ## <a name="news"></a> 🆕 News
 - **2023-10-27:** paper upload to arXiv.
-- **2023-12-24:** Github repo setup. Implementation code and pretrained model will be released after paper acceptance.
+- **2023-12-24:** Github repo setup. 
+- **2024-03-12:** Implementation code and pre-trained weight release
 
-## <a name="news"></a> 📖 Citation
+## <a name="usage"></a>  💻Usage
+ ### Prerequisites 
+ ```
+torch==1.5.1
+torchvision==0.6.1
+matplotlib==3.2.2
+pillow==8.0.1
+seaborn==0.10.1
+scikit-learn==1.0
+ ```
+ ### Pretrained model
+**[Setting 1](https://arxiv.org/abs/2310.18293)**, **[Setting 2](https://arxiv.org/abs/2310.18293)**
+ ### Testing
+To evaluate the model, you can pass --testing args to run different testing functions.
+```
+python test.py --test_dir $image_dir$ --out_dir $output_dir$ --gpu $gpu_id$
+```
+ ### Training
+ * Modify dataset path in ``dataset.py``
+ * Run the training code.
+ ```
+ python train.py --data_root $path_to_dataset_root$ --gpu $gpu_id$
+```
+**If there are any inconsistent to paper or coding bug, please feel free to raise an issue and let me know.** 
+
+## <a name="citation"></a> 📖 Citation
 If this work helps your research or work, please cite our paper
 ```
 @article{chen2023always,
@@ -22,9 +55,18 @@ If this work helps your research or work, please cite our paper
   year={2023}
 }
 ```
-## <a name="news"></a> 📜 License
+### 🙏 Acknowledgements
+Our code are based on [PyTorch-GAN](https://github.com/eriklindernoren/PyTorch-GAN),
+  and some useful functions from
+  [DLN](https://github.com/WangLiwen1994/DLN). We thank for contribution to the authors.
+
+### 📚 Our previous work
+Check out our previous work related to image restoration for adverse environment!
+* Domain adaptation for underwater image enhancement:  [UIESS](https://github.com/fordevoted/UIESS),
+
+## <a name="license"></a> 📜 License
 This project is released under [Apache 2.0 license](LICENSE)
 
 
-## <a name="news"></a> ✉️ Contact
+## <a name="contact"></a> ✉️ Contact
 Please free free let me know if you have any questions about this work by opening an issue or mail `210509fssh@gmail.com` `r09942066@ntu.edu.tw`

test.py

@@ -64,6 +64,7 @@ def findfolder(test_name, isTest=False):
                     help="name of the experiment")
 parser.add_argument("--test_dir", type=str, default=findfolder(test_name), help="path to test image directory")
 parser.add_argument("--out_dir", type=str, default=r'./output', help="path to output image directory")
+parser.add_argument("--testing", type=str, default=r'test_image', help="testing function. All list: ['test_image', 'test_plot_latent_tsne', 'test_latent_manipulation', 'eval_objective', 'test_QA']")
 parser.add_argument("--checkpoint", type=int, default=30, help="number of epoch of checkpoint")
 
 parser.add_argument("--n_downsample", type=int, default=2, help="number downsampling layers in encoder")
@@ -374,7 +375,7 @@ def align_to_four(img):
 
 if __name__ == '__main__':
 
-    testing = 'test_image'
+    testing = opt.testing
     testing_list = ['test_image', 'test_plot_latent_tsne', 'test_latent_manipulation', 'eval_objective',
                     'test_QA']