Two-Stage Single Image Reflection Removal
with Reflection-Aware Guidance

Due to the ill-posed nature, multiple images have been exploited to solve the reflection removal problem by providing extra information. Agrawal et al. [1] used two separate images as input, which are captured with and without flash light, respectively, and Fu et al. [5] further took into consideration the high frequency illumination. Schechner et al. [25] treated the input image as a superimposition of two layers, and novel blur kernels were proposed to eliminate mutual penetration of the separated layers, using images taken with different focal lengths as a clue. Polarization camera is also used to take advantage of the polarization properties of light [24, 36, 21]. In particular, based on the linear hypothesis $I=T+R$ in raw RGB space, Lei et al. [21] adopted a simple two-stage model but took multiple polarized images as input. In these methods, the behavioral difference between the reflection layer and the transmission layer under different settings serves as an effective clue for reflection removal.

Unlike the aforementioned methods, whose input images are captured in a fixed position, taking images from different viewpoints is another popular way to exploit multiple image information in reflection removal tasks [38, 31, 15, 8, 24, 6, 29, 40, 30, 9, 27]. In particular, Punnappurath et al. [20] leveraged the newly developed dual pixel (DP) cameras, and images captured from two sub-aperture views are readily accessible. Although approaches in the multi-input paradigm can obtain impressive results, the necessity to prepare qualified input images with specially designed hardware or extra human efforts makes them inconvenient, and the prerequisites are unable to be satisfied in many practical scenarios, which boosts the recent prevalence of single image reflection removal.

As the name implies, SIRR methods merely require a single image as input. Traditional methods employ various priors observed in real-world images to reduce the ill-posedness, such as gradient sparsity prior [13, 32, 2, 23], smoothness prior [16, 34] and ghosting cues [26]. Unfortunately, the representation and generalization ability of such methods is limited, and they are prone to produce poor results in the absence of discriminative hints, which is often the case when using handcrafted priors.

As a remedy, learning based methods have been proposed to solve the SIRR problem by leveraging the capacity of deep convolutional neural networks (CNNs), which mainly focus on data synthesis, loss function and network design. Due to the tremendous effort required to collect sufficient well-aligned real image pairs for model training, Fan et al. [4] generated reflection layers via smoothing and illumination intensity decay, which were then used to synthesize the observation by linear combination. Zhang et al. [41] further adapted the intensity decay to suit reflections with higher illuminance level, and applied random vignette to simulate different camera views. Besides, several loss functions have been particularly designed to leverage the characteristic of reflections. CEILNet [4] constrains the gradient maps of the predicted transmission layer to avoid blurry output. Zhang et al. [41] proposed an exclusion loss by minimizing the correlation between the gradient maps of estimated transmission and reflection layers, based on the observation that the edges in these two layers are unlikely to overlap. Li et al. [14] took advantage of the synthesis model to reconstruct the observation, which is constrained to approximate the input. As to network structure, most recent works follow two main schemes. On the one hand, [41] and [35] extracted input features via a pre-trained VGG-19 [28] model and stack them together for subsequent usage, which exploits multi-scale features and is tolerant of mild misalignment between real-world training pairs. On the other hand, the encoder-decoder framework is exploited in [4, 35, 39, 14], where skip connections are usually integrated for better feature utilization.

Considering the difficulty of directly estimating the transmission layer, cascaded models have been widely adopted in the hope of finding a better reflection removal pathway. Fan et al. [4] first predicted potential edge maps, and then the transmission layer is generated given the edge maps. Yang et al. [39] alternated between reflection and transmission estimation, while Li et al. [14] further incorporated a long short-term memory (LSTM) module [10] to generate the reflection and the transmission layer in a progressive manner. Such multi-stage methods are closely related to our proposed method, however, as further discussed in Sec. 3.1, they are still limited in exploiting the result from prior stage.

We note that the MIRR method [21] also uses a two-stage network but takes raw polarized images with different polarization angles as input. Moreover, it adopts the linear hypothesis $I=T+R$ and its concatenation-based fusion is limited in exploiting the estimated reflection for the subsequent stage. Empirically, the retraining of adjusted MIRR by feeding a single unpolarized image is still limited in removing heavy and bright reflections (See Fig. 1).

Given an observation $I$, CEILNet [4] predicts a potential edge map $E$ for the transmission layer, where a reflection smoothness prior is assumed to distinguish whether the image gradients belong to the transmission layer or the reflection layer. Unfortunately, such smoothness prior is often violated in real scenarios, and the resulting inaccurate edge maps will lead to a degraded performance by providing wrong information to the second stage.

BDN [39] and IBCLN [14] resort to predicting reflection layer in the prior stage. However, a simple concatenation is insufficient to make full use of the predicted reflection layer and is limited in suppressing the heavy and bright reflections, thus BDN [39] often fails to process the reflections in difficult regions (see Fig. 1). IBCLN [14] takes the input image $I$ and black image $\mathbf{0}_{h\times w}$ as initial estimation of transmission and reflection, and iteratively refines them with a LSTM [10] module to exploit cross-stage information, which significantly promotes the SIRR performance. Nevertheless, such principle makes the method less efficient and greatly relies on the generation quality in the prior stages, and the errors will also be accumulated across stages, leading to undesired artifacts in some scenarios (\eg, zooming in Fig. 1 for clear observation).

Considering the aforementioned limitations, we build a two-stage model, where a reflection-aware guidance (RAG) module is introduced in the second stage to better exploit the reflection layer predicted in the first stage. Specifically, we use a plain U-Net [22] model as the first stage subnetwork, namely $G_{R}$, which takes the observation $I$ as input and predicts the reflection layer, \ie, $\hat{R}=G_{R}(I)$. The structure of the second stage subnetwork $G_{T}$ is shown in Fig. 2, $\hat{R}$ is fed into $G_{T}$ together with the observation $I$ to generate the transmission layer, \ie, $\hat{T}=G_{T}(I,\hat{R})$. In particular, $G_{T}$ is composed of two encoders ($E_{I}$ and $E_{R}$) and one decoder ($D_{T}$), and an RAG module is elaborated in each block of $D_{T}$. Kindly refer to the supplementary material for detailed structure configurations.

As shown in Fig. 1, reflections are spatially nonuniform, whose location and intensity can be roughly predicted in the first stage. The estimated reflection $\hat{R}$ can then be incorporated with the observation $I$ for succeeding transmission estimation. However, the estimated reflection $\hat{R}$ may be inaccurate. Furthermore, the linear combination hypothesis may not hold true in regions with high reflection intensities, where few informative features of the transmission are retained in $I-\hat{R}$ (see Fig. 1). Therefore, estimating the transmission in such regions is more challenging, and is more likely to be an inpainting task. Taking these factors into account, we elaborate an RAG module in each decoder block for (i) generating feature complementary to the decoder and (ii) predicting a mask to indicate the heavy reflection regions for guiding image inpainting.

Given the observation feature $\mathbf{F}_{I}$, the reflection feature $\mathbf{F}_{R}$ and the decoder feature $\mathbf{F}_{\mathit{dec}}$, the RAG module is designed as shown in Fig. 2. Using $\mathbf{F}_{I}$ and $\mathbf{F}_{R}$, the difference feature $\mathbf{F}_{\mathit{diff}}$ is generated via a subtraction operation to better suppress the reflections, \ie, $\mathbf{F}_{\mathit{diff}}=\mathbf{F}_{I}-\mathbf{F}_{R}$, which serves as a complement and enhancement to the decoder feature $\mathbf{F}_{\mathit{dec}}$. Furthermore, taking $\mathbf{F}_{I}$, $\mathbf{F}_{R}$ and $\mathbf{F}_{\mathit{dec}}$ as input, the RAG module generates a mask $\mathbf{M}$ via two $1\times 1$ convolution layers followed by a sigmoid operation. Note that $\mathbf{M}=[\mathbf{M}_{\mathit{diff}},\mathbf{M}_{\mathit{dec}}]$, where $[\cdot,\cdot]$ denotes the concatenation operation, and $\mathbf{M}_{\mathit{diff}}$ and $\mathbf{M}_{\mathit{dec}}$ are masks regarding to $\mathbf{F}_{\mathit{diff}}$ and $\mathbf{F}_{\mathit{dec}}$, respectively. With such mask, the model is aware of the regions deviating from the linear combination hypothesis, and recovers the transmission layer in such areas mainly relying on the surrounding decoder feature using image inpainting. The mask $\mathbf{M}$ is visualized in Fig. 3.

Considering the similarity between inpainting and reflection removal in heavy reflection regions, we deploy a partial convolution [18] alongside each RAG module to better exploit the mask, which was first introduced in image inpainting methods [17, 37]. The mask generated by RAG module ranges from 0 to 1, thus the partial convolution in this paper can be formulated as,

where $\mathbf{W}$ and $\mathbf{b}$ represent the weight and bias of the partial convolution, $\mathbf{F}=[\mathbf{F}_{\mathit{diff}},\mathbf{F}_{\mathit{dec}}]$, $\ast$ and $\circ$ are convolution and entry-wise product respectively. $\mathbf{\bar{M}}_{3\times 3}$ contains the average values of $\mathbf{M}$ in $3\!\times\!3$ neighborhood regions, which can be efficiently calculated by a $3\!\times\!3$ average pooling operation.

Furthermore, we design a novel mask loss to better exploit the mask $\mathbf{M}$ during training RAGNet. As shown in Fig. 1, for areas with heavy reflections, $\mathbf{F}_{\mathit{diff}}$ can hardly provide informative features. Therefore, we dispose the values of $\mathbf{M}_{\mathit{diff}}$ to approximate $0$ in such areas, \ie,

where $i$ means the $i$-th layer, $\|\cdot\|_{{}_{1}}$ denotes the $\ell_{1}$ norm, $A[condition]$ represents the part of $A$ that meets the $condition$ in the square brackets, $\varphi$ is the threshold that delimits heavy reflection areas.

However, with $\mathcal{L}_{\mathrm{mask}}^{\mathit{diff}}$ only, $\mathbf{M}$ may be optimized towards a trivial solution $\mathbf{0}$, which diminishes the effect of $\mathbf{F}_{\mathit{diff}}$. Therefore, $\mathbf{M}$ is supposed to be $1$ for areas with few reflections, to avoid the trivial solution as a regularization term and force the partial convolution to exploit both $\mathbf{F}_{\mathit{diff}}$ and $\mathbf{F}_{\mathit{dec}}$, since both of them are reliable in such areas, \ie,

where $\xi$ means threshold for regions with few reflections.

Note that we do not constrain the mask values in remaining regions, which will be automatically optimized for better transmission estimation. Therefore, the mask loss is formulated as,

and we empirically set $\xi=0.01$ and $\varphi=0.3$. The effectiveness of mask loss will be further corroborated in ablation studies.

To train our RAGNet, several commonly used loss functions are collaborated, including reconstruction loss, perceptual loss [11], exclusion loss [41] and adversarial loss [7].

where $l$ indicates the index of $conv1\_2$, $conv2\_2$, $conv3\_2$, $conv4\_2$ and $conv5\_2$ layers. The weights $\{\kappa_{l}\}$ are used to balance different layers.

Taking the above loss functions into account, the learning objective to train our RAGNet can be formulated as,

where $\lambda_{1}=\lambda_{2}=\lambda_{5}=1,\lambda_{3}=0.2$ and $\lambda_{4}=0.01$.

We compare the proposed RAGNet with five state-of-the-art SIRR methods, \ie, CEILNet [4], Zhang et al. [41], BDN [39], ERRNet [35] , IBCLN [14], as well as the adjusted MIRR method by only taking the observation $I$ as the network input, \ie, Lei et al. [21]. We finetune these models on our training dataset and report the better result between the finetuned version and the released one. For Lei et al. [21], we adjust the network input and retrain the model with our training set.

The PSNR and SSIM metrics of all competing methods are reported in Table 2. Note that the results of Real 45 are omitted due to lack of ground-truth. Therefore, a user study is conducted on Real 45 to evaluate the reflection removal quality and the results also show the superiority of our RAGNet. The details are given in the supplementary material. It can be seen that the proposed RAGNet achieves the best performance by PSNR on three of the four datasets, and surpasses all other methods on average quantitative performance by a large margin ($\sim$0.8dB in PSNR). The quantitative gain against other deep cascaded models [4, 39, 14, 21] also indicates the benefit of appropriate reflection-aware guidance.

Fig. 4 shows the qualitative results and the corresponding ground-truth on three SIR² subsets and Real 20 dataset, and the results on Real 45 dataset are given in Fig. 5. One can see that CEILNet may fail when the reflections have clear edges, where the smoothness assumption is violated (\eg, the second row of Fig. 4). Zhang et al. [41] is prone to over-processing, which generates dim results with color aberration in most cases, while BDN [39] tends to generate brighter output images than input ones. Moreover, the reflections seem insufficiently removed in the results of ERRNet [35] and IBCLN [14]. In general, our method outperforms the competing methods, and removes the reflection areas more accurately and thoroughly. Please refer to the supplementary material for more qualitative results.

Our RAGNet is designed as a two-stage network, which takes reflection as an intermediate result to ease the reflection removal procedure. However, it may be concerned whether a well-designed one-stage network is sufficient for SIRR. To show the necessity and superiority of the two-stage architecture, we additionally design a single-stage variant, namely RAGNet_I→T, which directly estimates the transmission layer from the input.

According to Fig. 2, the generation of $T$ in RAGNet requires features from two sources, \ie, $\mathbf{M}_{\mathit{diff}}\circ\mathbf{F}_{\mathit{diff}}$ and $\mathbf{M}_{dec}\circ\mathbf{F}_{\mathit{dec}}$. We expect that $\mathbf{M}_{\mathit{diff}}\to 0$ when reflection is heavy as in Eqn. (2), so that only $\mathbf{F}_{\mathit{dec}}$ provides information for recovering heavy reflection areas. When designing RAGNet_I→T, we retain the functionality of the skip connection and $\mathbf{F}_{\mathit{dec}}$: one for providing auxiliary information via subtraction operations, the other focuses on gradually recovering $\hat{T}$. The structure of the one-stage model is provided in the supplementary material. The results in Table 3 and Fig. 6 show that the one-stage variant RAGNet_I→T suffers a huge performance degradation, indicating the essentiality of the two-stage architecture for our RAGNet.

Table 4 shows the quantitative results of the ablation studies on masks. It can be seen that our RAGNet surpasses its variants RAGNet_F and RAGNet_F∘M in both PSNR and SSIM, indicating the effectiveness of the deployment of mask and partial convolution. Furthermore, the ablation study on RAGNet${}_{\mathbf{F}\circ\mathbf{M}^{2c}}$ shows that predicting a per-channel mask better suits the SIRR task in our framework.

Fig. 7 shows the qualitative results of the ablation studies. It can be seen that, all of RAGNet_F, RAGNet_F∘M and RAGNet${}_{\mathbf{F}\circ\mathbf{M}^{2c}}$ fail to remove the heavy reflections, especially in complex environments (\eg, in the second row of Fig. 7, the light on the ground is very similar to the reflections). RAGNet outperforms the others by recovering the heavy reflection areas with contextual information, making the output image visually pleasant. Moreover, by integrating partial convolution, the reflection removal performance is also enhanced for images with mild reflections.

From the aspect of the proposed mask loss, we conduct three experiments, \ie, discarding $\mathcal{L}_{\mathrm{mask}_{\mathit{diff}}}$, $\mathcal{L}_{\mathrm{mask}_{\mathit{util}}}$ and both of them, respectively. Table 3 and Fig. 6 show that with the mask loss, the performance of our RAGNet is improved by a large margin. Furthermore, one can see that both items of the specifically designed mask loss are essential for RAGNet, and greatly boost the quantitative performance and visual quality. The thresholds $\xi$ and $\varphi$ in Eqns. (2) and (3) are discussed in the supplementary material.