S3Net: A Single Stream Structure for Depth Guided Image Relighting

Different from the conventional RGB tasks, extra depth information can improve accuracy in many computer vision tasks like semantic segmentation [17] and object detection [4]. Depth maps have demonstrated to be a useful cue to provide geometric and spatial information when combining with the RGB representation. In [4], they propose Faster-RCNN structure to tackle pedestrian detection. They prove that depth maps can be utilized to refine the convolutional features extracted from RGB images. Additionally, more accurate region proposals are achieved by exploring the perspective projection with the help of the depth information. In [17], a unified and efficient cross-modality guided encoder for semantic segmentation is proposed. This structure jointly filters and recalibrates both representations before cross-modality aggregation. Meanwhile, a bi-direction multi-step propagation strategy is introduced to effectively fuse information between the two modalities. Chen et.al [18] propose the approach to RGB-D salience object detection. The depth and image features are extracted by a dual backbone. Authors also integrate both features through densely connected structures and apply their mixed features to generate dynamic filters with receptive fields of different sizes. These works tend to use a dual backbone to extract features of different modalities. In practice, the size of relit images in VIDIT [16] and the corresponding depth map is very large (e.g., 1024$\times$1024). The relit images contain illuminant direction information, which is not proper to crop large images to small patches during training. To address this issue, we design a single stream structure to jointly extract both depth and the image features.

Following the rule in the 2021 NTIRE Depth Guided Image Relighting Challenge, there are two kinds of settings. First, the illuminant direction and the temperature are pre-defined [2, 6], which is known as one-to-one relighting. Second, the ambient condition is based on a guided image [3], which is known as any-to-any relighting. Both of them are very similar to other low-level vision tasks like image dehazing [14] image deraining [13], image smoke removal [19], image desnowing [20], reflection removal [21], and underwater image enhancement [12]. All of them belong to the image-to-image translation problems. To tackle those problems, the encoder-decoder network can be applied. Furthermore, when it comes to encoder-decoder structure, U-net [22, 23] is the most popular network for image-to-image translation tasks. This structure consists not only the encoder-decoder structure but also the skip connection, which concatenates the features with identical size from the encoder to the decoder. For example, Puthussery et.al [2] proposes a wavelet decomposed relit network for image relighting. This structure is a novel encoder-decoder network employing wavelet-based decomposition followed by convolution layers under a multi-resolution framework. Additionally, this network contains skip connection operations. Different from those low-level vision tasks, the relit image contains the information of illuminant direction and color temperature, which may cause the shadow of the objects to be located in different locations. In [3], a lighting-to-feature network is proposed to recover the corresponding implicit scene representation from the illumination settings, which is known as the inverse process of the lighting estimation network.

DWT [24] with orthogonal property is widely adopted in many computer vision tasks [23, 2]. DWT decomposes an image into various small patches in different frequency intervals, which can replace the existing down-sampling operations like max pooling or average-pooling. Therefore, many tasks apply DWT to diminish feature maps and achieve multi-scale features. Moreover, in [25], DWT is leveraged to design the objective function to measure the similarity between the ground truth and predict images. Motivated by these works, we also combine DWT in the loss function to make our network learn the multi-scale representations.

Attention mechanisms are important roles in both human perception systems and deep learning tasks [13, 26]. Attention mechanisms provide feature maps or certain sequence weights so that features of regions or locations can be magnified. Specifically, for computer vision tasks, attention mechanisms are categorized as the spatial attention [3] and the channel attention [13]. The former spatially utilizes weights to refine the feature maps, and the latter computes the global average-pooled features to implement the channel-wise attention. In this paper, both attention mechanisms are leveraged in our model to further increase the performance of the network.

The architecture of the proposed S3Net is presented in \figreffig:architecture. This network is based on [27] and contains the encoder and decoder parts. In the encoder part, we apply the Res2Net101 [28] as our backbone. The Res2Net can represent multi-scale features at a granular level and increases the range of receptive fields for each network layer. After the input is passed through the backbone, multi-scale feature extraction is achieved. Note that, the Res2Net utilized in this work discards the full connection layer and the size of final output feature maps from our encoder is $\frac{1}{16}$. The initial weight of the encoder is pre-trained parameters trained from the ImageNet. We connect bottom features to the decoder. The decoder consists of the stacks of convolution to refine the feature maps. Both the pixel shuffle [29] and the transposed convolution [30] are adopted to magnify feature maps. Furthermore, we leverage attention modules to refine intermediate features. Attention modules consist of residual layer [31], the spatial [26], and the channel attention [32] as shown in \figreffig:architecture (orange rectangle).

Furthermore, inspired by [11], we add the enhanced module in our S3Net. The enhanced module leverages average pooling in different strides to change the size of feature maps and receptive fields, which is effective to extract multi-scale features. Finally, the up-sampling is applied to restore diminished feature maps and all feature maps are concatenated. The enhanced module is illustrated in the light orange rectangle in \figreffig:architecture. Moreover, it is known that U-Net-like structure is beneficial in many tasks such as image dehazing [33, 14] and semantic segmentation [22]. Its skip connection encourages feature re-using. Therefore, in our S3Net, we also adopt skip connection to merge the last three feature maps from the backbone to their corresponding feature maps.

In practice, we directly combine the original image, original depth map, guided image and guided depth map as input. This input is seen as the 8-channel tensor and the output is the 3-channel relit image. In addition, we change the first convolution in the Res2Net backbone so that the S3Net can accommodate 8-channel tensors as input.

To train our S3Net, we apply three loss functions. The first function is Charbonnier loss [34] that can be regarded as the robust $L_{1}$ loss function. The Charbonnier loss is expressed as:

where $I$ and $\hat{I}$ represent the ground truth and relit images from the proposed network, respectively, and $e$ is seen as a tiny constant (e.g., $10^{-6}$) for stable and robust convergence. $L_{C}ha$ can restore global structure [34] and can be more robust to handle outliers.

Secondly, we apply the SSIM loss [35]. The SSIM loss is able to reconstruct local textures and details. It can be expressed as:

where $\sigma$ and $\mu$ represent the standard deviation, the covariance and the mean of images. In the image relighting task, to remove shadows from the original image, we extend the SSIM loss function so that our network can restore more detailed parts. We follow the method in [25] and combine DWT into the SSIM loss because lots of tasks [13, 23] have demonstrated that the DWT captures the high-frequency features, which is beneficial for reconstructing the clear details on relit images. Initially, the DWT decomposes the predicted image into four different and small sub-band images. The operation can be expressed as:

where superscripts mean the output from respective filters (e.g., $f_{LL}$, $f_{HL}$, $f_{LH}$ and $f_{HH}$).

$f_{HL}$ ,$f_{LH}$ and $f_{HH}$ are high-pass filters for the horizontal edge, the vertical edge and the corner detection, respectively. $f_{LL}$ is seen as the down-sampling operation. Moreover, the DWT can keep decomposing the $\hat{I}_{LL}$ to generate images with different scales and frequency information. This step is written as:

where the subscript $i$ means the output from the $i^{th}$ DWT iteration, and $\hat{I}^{LL}_{0}$ is the original predicted relit image. The SSIM loss terms described above are calculated from the original image pair and various sub-band image pairs. The fusion of the SSIM loss and the DWT is integrated as:

where $\gamma_{i}$ is based on [25] to control the importance of different patches.

The third loss is the perceptual loss [36]. Different from the aforementioned two loss functions, the perceptual loss leverages multi-scale features achieved from a pre-trained deep neural network (e.g., VGG19 [37]) to measure the visual feature difference between the ground truth and the estimated image. Formally, in this task, the VGG19 pre-trained on ImageNet is utilized as the loss function network. The perceptual loss is defined as

where $|\cdot|$ is the absolute value. The overall loss function is expressed as:

where $\lambda_{1}$, $\lambda_{2}$ and $\lambda_{3}$ are scaling coefficients and used to adjust the relative weights on the three components.

The dataset used in the 2021 NTIRE image challenge of depth guided image relighting is the Virtual Image Dataset for Illumination Transfer (VIDIT) [16]. This dataset consists of 390 different scenes which are with 40 different illumination settings (five different color temperatures from 2500 to 6500K and 8 azimuthal angles). There are 15600 images totally. Additionally, the corresponding 390 depth maps are provided. The size of all training images and depth maps are 1024 × 1024 × 3 and 1024 × 1024 × 1, respectively. 300 virtual scenes with different illumination settings are applied for training while the other 90 virtual scenes are used for validation and testing.

In the training phase, we first choose an arbitrary image and the corresponding depth map. Then, we choose an image as a guided image that contains random color temperature and azimuthal angles. With this information, the output image can be determined. Our network utilizes the original image, the original depth map, the guided image and the guided depth map as input to reconstruct the final decided image. We do not apply any data augmentation like random flip and random cropping during the training phase.

we apply the AdamW [38] as the new optimizer, and the batch size is set as 3 and the network is trained for 100 epochs with the momentum $\beta_{1}$ = 0.5 and $\beta_{2}$ = 0.999. The learning rate is set as $10^{-4}$ and it is divided by ten after 20 epochs. The weights of these three loss terms (i.e., $\lambda_{1},\lambda_{2}$ and $\lambda_{3}$) are set as 1, 1.1 and 0.1, respectively. All experiments are performed on the PyTorch platform and a single Nvidia V100 graphic card. We spend about 80 hours finishing the model training.¹¹1The source code will be released in our project page.

To find the best effectiveness of the proposed S3Net, some ablation experiments are conducted in this section. The peak signal-to-noise ratio (PSNR) and the structural similarity (SSIM) are adopted as objective metrics for quantitative evaluation.

The ablation experiments consist of three experimental settings: 1) The input is only with the source image and guided image but without depth maps, that is, the input is a six-channel tensor (Image); 2) The input is with the source image, guided image, and their corresponding depth maps (Both); For these two settings, $L_{Cha}$, $L_{SSIM}$ and $L_{Per}$ are adopted as objective functions. 3) The same setting in 2) but the $L_{SSIM}$ [25] is replaced by $L_{W-SSIM}$. The results are reported in Table 1. One can see that the PSNR and SSIM scores of setting 2 can be improved compared with setting 1. It can prove that using depth maps can improve the performance of relighting because they can provide more spatial information for the network. Moreover, compared with setting 2, the performance of setting 3 is improved effectively. It indicates that the $L_{W-SSIM}$ in the network can be more beneficial compared with the $L_{SSIM}$.

We also present some visual comparisons from VIDIT validation set for depth guided any-to-any relighting problem to prove the effectiveness of using the wavelet SSIM loss in the proposed S3Net. As shown in \figreffig:vis, images relit from the model trained with wavelet SSIM loss are closer to the ground truth. This model can properly remove shadows and transfer the ambient conditions from guide images to source images.

We list the results of the proposed S3Net compared with other competing entries in depth guided any-to-any relighting challenge of NTIRE 2021 workshop [10] in Table 2. Besides PSNR and SSIM, the Mean Perceptual Score (MPS) defined as the average of the normalized SSIM and LPIPS [39] scores, themselves averaged across the entire test set of each submission is adopted to evaluate performance of all submissions. As shown in Table 2, our results obtained the $3^{rd}$ place, the $4^{th}$ place, the $2^{nd}$ place and the $3^{rd}$ place in terms of SSIM, PSNR, LPIPS and MPS. It is noted that our model takes 2.042 seconds on average to generate a relit result with 1024 $\times$ 1024 image size in the test phase.

The proposed method learns the mapping functions of depth guided image relighting. Although our method achieves competitive performance in this competition, it may fail under the certain conditions. As shown in \figreffig:fail, the relit images are very different from the ground truth. Because when original images contain a large region of shadows, our model cannot identify the foreground and the backward of them. Even the depth maps are given, this information just provides the front side instead of the omnidirectional spatial information. Therefore, the color temperature from the guided images is transferred to the relit image (e.g., \figreffig:fail (e)), but the model reconstructs the poor structures.