SAPNet

This repository contains the official Pytorch implementation of the paper: "SAPNet: Segmentation-Aware Progressive Network for Perceptual Contrastive Image Deraining"

Updates:

Code will be updated before 2021/11/23 **Arxiv Link is available at https://arxiv.org/abs/2111.08892

Abstract

Deep learning algorithms have recently achieved promising deraining performances on both the natural and synthetic rainy datasets. As an essential low-level pre-processing stage, a deraining network should clear the rain streaks and preserve the fine semantic details. However, most existing methods only consider low-level image restoration. That limits their performances at high-level tasks requiring precise semantic information. To address this issue, in this paper, we present a segmentation-aware progressive network (SAPNet) based upon contrastive learning for single image deraining. We start our method with a lightweight derain network formed with progressive dilated units (PDU). The PDU can significantly expand the receptive field and characterize multi-scale rain streaks without the heavy computation on multi-scale images. A fundamental aspect of this work is an unsupervised background segmentation (UBS) network initialized with ImageNet and Gaussian weights. The UBS can faithfully preserve an image's semantic information and improve the generalization ability to unseen photos. Furthermore, we introduce a perceptual contrastive loss (PCL) and a learned perceptual image similarity loss (LPISL) to regulate model learning. By exploiting the rainy image and groundtruth as the negative and the positive sample in the VGG-16 latent space, we bridge the fine semantic details between the derained image and the groundtruth in a fully constrained manner. Comprehensive experiments on synthetic and real-world rainy images show our model surpasses top-performing methods and aids object detection and semantic segmentation with considerable efficacy.

Preparing Dataset

First, download training and testing dataset from either link BaiduYun OneDrive

Next, create new folders called dataset. Then create sub-folders called train and test under that folder. Finally, place the unzipped folders into ./datasets/train/ (training data) and ./datasets/test/ (testing data)

Training

Run the following script in terminal

python train.py

Testing

Run the following script in terminal

bash main.sh

Hyperparameters

General Hyperparameters

Train/Test Hypeparameters

Contact

Please reach [email protected] for further questions. You can also open an issue (prefered) or a pull request in this Github repository

Acknowledgement

This repository is borrowed heavily from PreNet. Thanks for sharing!

TODO List

• Upload Pretrained Weight
• Add Visual Comparisons
• Add References
• Upload Arxiv Link
• Upload BibTeX