Unsupervised Raindrop Removal from a Single Image using Conditional Diffusion Models

Raindrop removal has seen several applications, such as in unmanned aerial vehicle inspection [7] and coastal video enhancement [8]. Most approaches toward raindrop removal rely on GANs [9, 10, 11, 12, 13, 7]. Most recently, Xia et al. [9] devised a two-step GAN to maintain balance between raindrop removal and image inpainting. Their strategy is quite similar to our own but differs in the use of architecture. In contrast to GANs, we employ diffusion models to perform the inpainting. Recently, Özdenizci and Legenstein [14] designed TransWeather, a patch-based diffusion model for vision restoration in adverse weather conditions. Unlike this study, they approach multiple weather conditions at once, including rain, snow, and haze. Although we use a similar diffusion approach, we are strictly focused on raindrop-filled images.

The theory of diffusion models is developed in detail by Ho et al. [15]. A diffusion model is essentially a Markov chain containing $T$ latent variables $x_{1},\dots,x_{T}$ with respect to the original image data $x_{0}$. The model consists of (1) a forward process, which gradually adds noise to $x_{0}$ and returns the latent variables in succession, and (2) a reverse process for denoising, which gradually removes noise from the latent variables to eventually reobtain $x_{0}$. The forward process is represented as distributions $q(x_{t})$, and the reverse process is represented as distributions $p_{\theta}(x_{t})$. The illustrating diagram is provided in Figure 2.

For the forward process, the conditional probability of $x_{t}$ given $x_{t-1}$ can be calculated using the following formula:

where $\beta_{t}$ is a variance hyperparameter for the forward process. Let $\alpha_{t}=1-\beta_{t}$ and $\overline{\alpha}_{t}=\prod_{s=1}^{t}\alpha_{s}$. The conditional probability of $x_{t}$ given $x_{0}$ can be directly calculated using the following formula:

For the reverse process, the conditional probability of $x_{t-1}$ given $x_{t}$ can be calculated using the following formula:

Diffusion models have numerous applications to image synthesis [16, 17, 18], in which they have become the state-of-the-art approach. In addition, they have been applied to anomaly detection [19] and image restoration [14]. This indicates the versatility and effectiveness of diffusion models in the field of image processing.

A raindrop mask is a binary image that maps areas containing raindrops. This map separates the raindrop areas and the background areas of the image. If an area contains raindrops, then the pixels in the are set to $255$ (white). For the background areas that do not contain raindrops, pixels are set to $0$ (black).

We consider two approaches for generating a pseudo raindrop mask: residual mask generation and synthetic mask generation. In the first approach, we create a residual mask simply by subtracting the clean image from the corresponding image with raindrops. Then, we set a threshold to extract the raindrop regions from the resulting residual image. Residual pixels that lie above the threshold are categorized as part of the raindrop area, whereas those that do not are categorized as part of the background. Unfortunately, directly calculating the residual on images often results in poor masking due to the existence of extraneous variables. These variables include brightness, light intensity, other objects, etc. So, we employ several pre-processing techniques to minimize the influence of these variables, including histogram equalization and photometric distortion. In the second approach, we employ a more sophisticated method to generate synthetic masks. We then train a supervised raindrop detector by supplying it with a synthetic raindrop ground truth.

Given a pseudo-mask $m$ generated from the previous module, we aim to restore the background on the masked areas. In this module, we employ a pre-trained DDPM model [4] that has been trained on clean images of the Raindrop dataset. Specifically, we run the inference process using the pre-trained model to restore the background. We expect that the model can learn the distribution of clean images so it can generate novel images from that distribution. The intuition is simple. If the model can generate such novel images, then it can be used to reconstruct the background image when conditioned on the raindrop mask.

Background reconstruction based on a mask essentially consists of predicting the missing pixels due to raindrops conditioned on the mask. The Raindrop dataset comes with ground-truth clean images, so we can evaluate the result using these ground truths. First, we denote $x$ as the ground-truth image, $m\odot x$ as the missing pixels, and $(1-m)\odot x$ as the known pixels. The operator $\odot$ denotes element-wise multiplication.

Recall that from the DDPM perspective, model training is a forward diffusion process, while model inference is a reverse diffusion process. From Equation 3, we can see that in a single reverse step, we only need the variable $x_{t}$ to reconstruct $x_{t-1}$. We can ensure that DDPM constructs only on the masked regions via a simple trick. After we obtain $x_{t-1}$, we can alter the known pixels in $(1-m)\odot x_{t-1}$ using the corresponding pixels from $x_{t-1}$ that are obtained by forward diffusion. Fortunately, we can directly go to $x_{t}$ from $x_{0}$ for any $t$ by using Equation 2. So, we conduct the background restoration by combining two directions: forward and reverse. We use Equation 3 to obtain the unknown pixels and Equation 2 to obtain the known pixels. This technique can be formally defined as follows:

where $x_{t-1}^{\text{known}}\sim\mathcal{N}\left(\sqrt{\overline{\alpha}}x_{0},(1-\overline{\alpha}_{t})\mathbf{I}\right)$ and $x_{t-1}^{\text{unknown}}\sim\mathcal{N}\left(\mu_{\theta}(x_{t},t),\Sigma_{\theta}(x_{t},t)\right)$. Figure 5 illustrates the process of background restoration.

In the experiments, we use two datasets: the Raindrop dataset by Qian et al. [2] and the Cityscapes dataset [6]. The former dataset, Raindrop, contains images with raindrops and their corresponding clean images as ground truth. Our main goal is to detect and remove raindrops from these images and evaluate the results against clean images. This dataset consists of $861$ train data, $58$ validation data, and $249$ test data. Each image has a spatial resolution of $720\times 480$ with RGB channels. Most images in this dataset are photos taken around campus.

The latter dataset is the Cityscapes dataset. We use the gtFine_trainvaltest.zip package contained in this dataset. This package contains $5000$ images, along with parameter settings on the camera that is used to take the photos. We employ this dataset to train our model. We generate synthetic raindrops on this dataset and then use the result to train the raindrop detector.

We use 1 GPU V100 for all experiments. We implement the raindrops synthetic generator using the C++ programming language, adapting the code from Hao et al. [27]. To better handle the image data using C++, we use OpenCV, an open-source C++ library that includes several hundreds of computer vision algorithms.

We implement the raindrop detector using the Python 3.8 programming language and the PyTorch 1.12 library to better create the model architecture. For easier model training and fine-tuning, we use the PyTorch Lightning framework [28]. We train the model for $100$ epochs with batch size $32$. We also use the AdamW optimizer [29] with learning rate $10^{-3}$ and weight decay $10^{-4}$. For the learning rate scheduler, we use StepLR with step_size $5$. We train the model to minimize the binary cross entropy loss as the raindrop detector task is essentially a segmentation task utilizing per-pixel classification.

For the background reconstructor, we use the DDPM model implemented by Dhariwal et al. [4]. We train the model on the clean images of the Raindrop dataset for 1M iterations. We set the number of diffusion steps to $T=1000$. To save memory, we center-crop the images to size $128\times 128$. We use the AdamW optimizer with learning rate $10^{-4}$ without weight decay. The training process took 6 days in total.