Learning Scene Structure Guidance via Cross-Task Knowledge Transfer for Single Depth Super-Resolution

Figure 2 shows the overall architecture of our framework, which mainly consists of three components: depth super-resolution network (DSRNet), depth estimation network (DENet) and the middle cross-task interaction module. Given a collection of paired LR-HR depth maps $\{D_{lr}^{(k)},D_{hr}^{(k)}\}_{k=1}^{M}$ with the corresponding HR color images $\{I^{(k)}\}_{k=1}^{M}$ as training data, where $M$ is the number of training data, our goal is to learn a model, i.e., DSRNet, that can predict the super-resolved depth maps $\{D_{sr}^{(k)}\}_{k=1}^{M}$ from their corresponding downsampled versions $\{D_{lr}^{(k)}\}_{k=1}^{M}$.

Specifically, the structure of DSRNet is designed based on a network unit from deep back-projection network (DBPN) [14], which can effectively improve the feature representations through iterative projecting HR representations to LR spatial domain and then mapping back into HR domain¹¹1We direct readers to refer to [14] for more details about the design of DBPN block. . The shallow features $F_{sr}^{shallow}$ are first extracted from $D_{lr}$ through a simple convolutional block (including three convolutional layers), and then sent into $N$ stacked DBPN blocks to obtain HR features $\{F_{sr}^{n}\}_{n=1}^{N}$. $D_{sr}$ is finally reconstructed from $F_{sr}^{N}$ through another convolutional block. DENet takes $I$ as input and estimates the depth map $D_{de}$. The architecture of DENet is similar to DSRNet, but replaces the DBPN blocks with deeper residual blocks  [15] to extract informative features $\{F_{de}^{n}\}_{n=1}^{N}$ from color image.

The cross-task interaction module acts as a bridge to connect DSRNet and DENet, and realizes bilateral knowledge transfer between them. It consists of two components, i.e., a cross-task distillation scheme and a structure prediction network (SPNet), where the former focuses on the interaction between multi-scale features extracted from both networks while the latter uses structure maps as supervision to further guide the learning of both networks.

Note that, at the training stage, DSRNet, DENet and the cross-task interaction module are jointly learned by using both color image and depth map as input. In testing phase, DSRNet is the final choice to predict HR depth map from only LR depth map without the help of color image.

Knowledge distillation is generally viewed as a technique of transferring beneficial information from a top-performing model to the other naive one. Different from the commonly used distillation techniques, in which the teacher network is trained beforehand and fixed under the assumption that it always learns a better representation than the student network, our goal is to train SRNet and DENet collaboratively and encourage them to benefit from each other.

Inspired by mutual learning methods [52, 22], we propose a cross-task distillation scheme, in which the roles of teacher and student will exchange between both tasks based on their current performances on depth recovery in the iterative collaborative training. Specially, at the current round of training, we need to determine the teacher in advance according to their performance at the previous round. We compute the average pixel error between each recovered depth map and its groundtruth for both networks:

where $\{H,W\}$ are the size of the output depth map. If $e_{dsr}$ is smaller than $e_{de}$, DSRNet has a relatively better performance, and becomes the dominant one to guide the learning of DENet, and vice versa.

Next, in order to distill more meaningful knowledge that can accurately describe the essential structural characteristics of depth map, we introduce a multi-space distillation scheme to condense the knowledge from the perspectives of output and affinity spaces, as shown in Figure 2.

Assuming the dimension of feature $F$ is $w\times h\times c$, the reshape function $\mathbb{R}$ recasts $F$ as $\mathbb{R}(F)$ with the dimension of $wh\times c$. The affinity matrix $A$ is defined as:

where $\sigma(\cdot)$ is the softmax operation, $\otimes$ is the matrix multiplication and $T$ is the transpose operator. The distillation loss of affinity space is defined as the following,

The final distillation loss $\mathcal{L}_{distill}$ is expressed as:

where $\gamma$ is an adjustment parameter. Note that, $\mathcal{L}_{distill}$ should be imposed on the training of the student, but not the teacher, which are determined by the errors comparison between $e_{dsr}$ and $e_{de}$.

The goal of SPNet is to predict a structure map $S$ from the last feature maps $F_{sr}^{N}$ and $F_{de}^{N}$ generated by DSRNet and DENet respectively. Through the supervision with the groundtruth structure map $S_{gt}$²²2$S_{gt}$ is obtained by computing the difference between adjacent pixels (gradients) in $D_{HR}$. Following [30], we calculate it by a convolution layer with a fixed kernel to extract the high-frequency parts from $D_{HR}$., SPNet can provide extra structure regularization to help both DSRNet and DENet learn more informative structure representations to alleviate the problem of RGB-D structure inconsistency. As shown in Figure 3, SPNet consists of a fusion module and a structure CNN, where the latter is a lightweight network with five stacked ‘Conv+ReLU’ layers and a last ’Conv’ layer.

Usually, the erroneous recovery of DSR and DE tasks occurs in the regions around depth boundaries and fine structures in a depth map, which are subject to higher recovery uncertainty. Therefore, instead of simply concatenating $F_{sr}^{N}$ and $F_{de}^{N}$ and sending into the structure CNN, we design an uncertainty-induced attention fusion module to strengthen these informative structure features by attending the recovery uncertainty to the feature map. Thus, we first compute the uncertainty maps $U_{sr}$ and $U_{de}$ of both networks by activating the recovery errors:

where $Conv_{1\times 1}$ is the $1\times 1$ convolution to adjust the channels. Then, we use the uncertainty maps to re-weight $F_{sr}^{N}$, $F_{de}^{N}$ and fuse them through an attention module:

where $F_{struc}$ is the fused features for structure prediction. [$\cdot$] denotes concatenation operation and $*$ is element-wise multiplication. Note that, through back-propagating the network gradients from SPNet in the backward information flow, the parameters of DSRNet and DENet can be updated.

The training process of our framework can be divided into two steps, as presented in Algorithm 1. First, we separately train DSRNet and DENet with the groundtruth $D_{hr}$. The losses are defined as follows:

where $\mathcal{L}_{DSR}$ is a common pixel-wise L1 loss for the task of DSR. Following [10], $\mathcal{L}_{DE}$ is set as a combination of the reconstruction loss (L1 loss) and structural similarity (SSIM). $\lambda$ is an adjustment parameter.

Then, we introduce the cross-task distillation between both networks with the loss $\mathcal{L}_{distill}$ in Eq. (6). At the same time, we randomly initialize SPNet, and train it together with DSRNet and DENet. The loss for SPNet is defined as:

where $\mathbb{G}(\cdot)$ denotes SPNet, $F_{struc}$ is the fused structure feature in Eq.(9) and $S_{gt}$ is the ground truth of the structure. If DSRNet is chosen as the student, the parameters of DENet are fixed at the current epoch, and DSRNet is updated with the following loss:

where $\rho_{1},\rho_{2}$ are the trade-off parameters. Otherwise, DSRNet is fixed, and DENet is updated with the loss $\mathcal{L}$ by replacing $\mathcal{L}_{DSR}$ with $\mathcal{L}_{DE}$.

Middlebury Dataset. We compare our method with state-of-the-art (SOTA) DSR methods on Middlebury under four different up-scaling factors ($\times 2$, $\times 4$, $\times 8$ and $\times 16$). All the compared methods and ours are deep learning based methods, which are trained and tested under the same conditions for fair comparison. Table 1 summarizes the quantitative results on the Middlebury dataset. ‘DSRNet w/o CT’ and ‘DSRNet’ denote the results of the proposed method without and with cross-task interaction scheme, respectively. Benefiting from the backbone of DBPN network, the performance of ‘DSRNet w/o CT’ goes beyond most of the previous methods, but slightly inferior to the recent SOTA methods, e.g., DSRN. When training together with DENet assisted by cross-task interaction, the performance of our DSRNet steps further compared to ‘DSRNet w/o CT’, and is comparable to DSRN, even better than it on at least half of the cases. Note that, different from these color-guided methods, we stand on single DSR without the help of color image in testing phase, but achieves satisfactory results on both accuracy and runtime (evaluated later in this section). The $\times$8 DSR results are visually shown in Figure 4. For the structural details, e.g., the stick and the teapot handle in Art, and the toy’s contour in Dolls, we clearly recover these regions without introducing the texture copying artifacts thanks to our structure prediction network (SPNet).

NYU Dataset. We further evaluate our method on NYUv2 dataset compared with other SOTA methods in Table 2. Our method performs the best on all the cases with different up-scaling factors. Figure 5 presents visual comparisons under the $\times$16 DSR case. Observing from the recovered depth maps and enlarged patches, our method achieves the clearest results and preserves the sharpest and correct object contours, which demonstrates our effectiveness to learn the information of scene structure guidance from color image.

ToFMark Dataset. We evaluate the generalization ability of our method on ToFMark dataset. Following DGDIE [11], we fill the missing points in depth maps and downsample them by $\times$2 downsampling factors, then send them to the $\times$2 model trained on Middlebury dataset to acquire the generalization results. As shown in Table 3, our method obtains the best objective generalization results for all three test examples. Figure 6 shows the visual comparison on Shark. Our method has higher visual quality and less blur than others, especially at the boundary regions.

Runtime. In Figure 8, we summarize the overall performance by the tradeoff between accuracy and runtime. We measure the runtime for $\times$8 DSR to their full resolution (about 1080$\times$1320) on Middlebury dataset. The color-guided methods i.e., JGF [29] MSG [18], and DSRN [40], run slower than ours due to the usage of HR color images in testing phase. MS is the single DSR version of MSG, but still inferior to ours for both accuracy and speed. Owing to the cross-task interaction, we achieve satisfactory recovered results with minimum inference time.

In this section, we further verify the key designs of our cross-task interaction, i.e., cross-task distillation and structure prediction task by ablation study.

We report the $\times$8 DSR results on Middlebury and NYUv2 dataset under different experimental settings, as shown in Table 4. The $1^{st}$ and $2^{nd}$ rows are the objective results of initial DSRNet and DENet trained separately without cross-task interaction scheme. DENet performs relatively inferior to DSRNet, because estimating depth map from only color image is more difficult than from the degraded LR version. Next, $\mathcal{L}_{S}$ and $\mathcal{L}_{A}$ distill knowledge from output space and affinity space respectively. When gradually integrating each loss into the baseline, it brings a progressive performance improvement (6.5% and 8.9% increase on Middlebury and NYU respectively), which verifies the effectiveness of cross-task distillation strategy is remarkable. Besides, we also visualize the error maps between recovered depth map and groundtruth to validate the process of our iterative collaborative training in Figure 9. As the epoch goes, the recovery errors for both DSRNet and DENet are decreased at the boundary regions. Finally, after introducing the SPNet (the final case in Table 4 ), the values of both error metrics are further decreased (4.3% error reduction against the previous case). Figure 9(e) also validates this from visual performance.

Besides, we conduct ablation study to evaluate the effectiveness of our uncertainty-induced attention fusion module as presented in Table 5. Compared to simply concatenating features in channel dimension and sending into SPNet (SPNet w/o U), our proposed fusion module can offer significant assistance for SPNet to predict the structure map. We also visualize the feature maps in the fusion process in Figure 7. The object boundaries, e.g., the regions around the toy head and the strip behind, are obviously highlighted by the uncertainty map, which facilitate the final structure perception of SPNet.