S³M-Net: Joint Learning of Semantic Segmentation and Stereo Matching for Autonomous Driving

Semantic segmentation has been a long-standing problem in the field of computer vision over the past decade [6, 29]. The SoTA networks in this research area can generally be classified into two categories: (1) single-modal networks with a single encoder and (2) feature-fusion networks with multiple encoders [3, 30, 31]. In the early attempts to tackle semantic segmentation, researchers primarily focused on encoder-decoder architectures for pixel-level classification. Notable examples include SegNet [13], U-Net [14], PSPNet [15], the DeepLab series [16, 17], and Transformer-based networks [32, 33, 34]. The encoder extracts hierarchical deep features from the input image, while the decoder produces the segmentation map by upsampling and combining the features from different encoder layers. However, these networks are limited in their ability to effectively combine deep features extracted from different modalities (or sources) of visual information. As a result, they often struggle to produce accurate segmentation results in challenging scenarios characterized by poor lighting and illumination conditions [3]. Therefore, researchers have turned their focus towards feature-fusion networks that can effectively integrate deep features learned from multiple modalities (or sources) of visual information. This problem is commonly referred to as “RGB-X semantic segmentation”, where “X” represents the additional modality (or source) of visual information, in addition to the RGB images. The most representative feature-fusion networks based on convolutional neural networks (CNNs) include FuseNet [19], MFNet [35], RTFNet [20], and our previous works SNE-RoadSeg series [3, 21]. Furthermore, Transformer-based RGB-X semantic segmentation networks, such as OFF-Net [22] and RoadFormer [24], have been recently introduced. In this article, we design our S³M-Net based on the SNE-RoadSeg architecture and explore more effective solutions for the feature fusion operation.

Conventional explicit programming-based stereo matching algorithms (local, global, and semi-global) generally consist of four main procedures: (1) cost computation, (2) cost aggregation, (3) disparity optimization, and (4) disparity refinement [26]. The performance of these algorithms has been significantly outperformed by end-to-end deep stereo networks, thanks to the recent advancements in deep learning techniques. PSMNet [36], GwcNet [37], AANet [38], LEA-Stereo [39], RAFT-Stereo [40], and CRE-Stereo [41] are six representative end-to-end deep stereo networks proposed in recent years. PSMNet [36] employs a spatial pyramid to capture multi-scale information and employs multiple 3D convolutional layers to exploit both local and global contexts for cost computation. GwcNet [37], on the other hand, builds upon the foundation of PSMNet by constructing the cost volume via group-wise correlation, thereby enhancing the 3D stacked hourglass network. In light of the computational demands of 3D convolutions, researchers have actively sought ways to minimize the trade-off between efficiency and accuracy in stereo matching. For example, LEA-Stereo [39] introduces the neural architecture search (NAS) [42] technique to stereo matching. This pioneering approach results in the first end-to-end hierarchical NAS framework for deep stereo matching. RAFT-Stereo [40], a rectified stereo matching method that draws inspiration from the optical flow estimation network RAFT [43], leverages the RAFT architecture to perform accurate and real-time stereo matching inference. The network utilizes recurrent structures to refine correlation features and enhance the disparity estimation accuracy. CRE-Stereo [41], another recent prior art based on recurrent refinement (to update disparities in a coarse-to-fine manner) and adaptive group correlation (to mitigate the impact of erroneous rectification), yields more compelling disparity estimation results. In this article, we develop our S³M-Net based on the RAFT-Stereo architecture.

Existing frameworks that jointly address semantic segmentation and stereo matching generally focus on improving disparity accuracy by leveraging semantic information [7, 8, 9, 11, 12, 10], while the discussion regarding the utilization of disparity information to enhance semantic segmentation at the feature level for joint learning remains limited, except for the exploration of “RGB-X semantic segmentation” discussed in Sect. II-A. Nevertheless, these prior arts either require a large amount of well-annotated training data or involve intricate training strategies for the joint learning of both tasks. For instance, SegStereo [7] and DispSegNet [8] require an initial unsupervised training phase on the large-scale Cityscapes [44] dataset, followed by a subsequent supervised fine-tuning on the smaller KITTI 2012 and 2015 [45, 28] datasets. Similarly, the studies presented in [9, 11, 12] involve the pre-training of their spatial branches (performing stereo matching) on the large-scale SceneFlow [46] dataset, followed by the fine-tuning of both semantic and spatial branches on the KITTI 2012 and 2015 datasets [45, 28]. DSNet [10] adopts a different joint learning strategy in which the training alternates between the semantic segmentation and stereo matching networks, with each network being frozen during the training of the other. However, achieving simultaneous convergence of the two networks can be challenging, as the shared features are not learned in an end-to-end manner. Additionally, we were unable to locate publicly available source code (in PyTorch or TensorFlow) for these prior arts, and re-implementations carry the risk of introducing errors. In contrast to the aforementioned approaches, our proposed S³M-Net is trained in an end-to-end fashion and capable of jointly learning semantic segmentation and stereo matching even when the training data are limited.

Given a pair of well-rectified stereo images $\boldsymbol{I}^{L}$, $\boldsymbol{I}^{R}\in\mathbb{R}^{H\times W\times 3}$, where $H$ and $W$ denote their height and width, respectively, we employ a joint encoder consisting of a series of residual blocks and downsampling layers to extract features $\mathcal{F}^{L}=\left\{\boldsymbol{F}^{L}_{1},\dots,\boldsymbol{F}^{L}_{n}\right\}$ and $\mathcal{F}^{R}=\left\{\boldsymbol{F}^{R}_{1},\dots,\boldsymbol{F}^{R}_{n}\right\}$ from $\boldsymbol{I}^{L}$ and $\boldsymbol{I}^{R}$, respectively. $\mathcal{F}^{L}$ is subsequently shared with the semantic segmentation task.

Using the features $\mathcal{F}^{L}$ and $\mathcal{F}^{R}$ extracted by the joint encoder, we first construct an initial 3D correlation volume $\boldsymbol{C}_{1}\in\mathbb{R}^{H\times W\times W}$ as follows:

where $i$ represents the $i$-th row, and $j$ and $k$ represent to the $j$-th and $k$-th columns in the left and right shared feature maps, respectively. We then construct a pyramid of 3D correlation volumes $\mathcal{C}=\left\{\boldsymbol{C}_{1},\dots,\boldsymbol{C}_{m}\right\}$ by downsampling $\boldsymbol{C}_{1}$ with average pooling operations. The $m$-th 3D correlation volume $\boldsymbol{C}_{m}\in\mathbb{R}^{H\times W\times\frac{W}{2^{m-1}}}$ is constructed from the ($m-1$)-th 3D correlation volume $\boldsymbol{C}_{m-1}$ using 1D average pooling with a kernel size of 2 and a stride of 2. Inspired by RAFT-Stereo [40], we adopt a multi-level GRU update operator to refine a sequence of disparity maps $\mathcal{D}=\left\{\boldsymbol{D}_{1},\dots,\boldsymbol{D}_{n}\right\}$, where $\boldsymbol{D}_{i}\in\mathbb{R}^{H\times W}$ ($i=1,\dots,n$). This refinement process is performed in a coarse-to-fine manner, starting from an initial disparity map $\boldsymbol{D}_{0}$ in which all disparities are initialized to $0$.

In stereo matching, a lower number of channels, e.g., the 256 channels utilized in RAFT-Stereo [40], is often sufficient for capturing relevant features for 1D correspondence search, especially when considering computational efficiency. On the other hand, semantic segmentation requires pixel-level classification and a more in-depth scene understanding. It benefits from complex feature representations that can capture fine-grained details and object boundaries, making a larger number of channels, e.g., the 2048 channels employed in SNE-RoadSeg [3], advantageous for this task. Therefore, we introduce the FFA module to align the channels and resolutions between the disparity and semantic feature maps during joint learning.

As illustrated in Fig 2, given the left shared feature maps $\mathcal{F}^{L}=\left\{\boldsymbol{F}^{L}_{1},\dots,\boldsymbol{F}^{L}_{n}\right\}$ and the disparity map pyramid $\mathcal{D}=\left\{\boldsymbol{D}_{1},\dots,\boldsymbol{D}_{n}\right\}$, we obtain the adapted fused feature sequence $\mathcal{F}^{F}=\left\{\boldsymbol{F}^{F}_{1},\dots,\boldsymbol{F}^{F}_{n}\right\}$ using our proposed FFA module, which can be formulated as follows:

where $\mathcal{E}^{D}$ denotes the disparity map encoding operation, $\oplus$ denotes the feature fusion operation, and $\mathcal{A}_{i}$ is defined as our feature adaptation operation, as formulated as follows:

where $\mathcal{R}$ represents the remapping operation from the shared feature space to the semantic feature space, and $\mathcal{E}$ represents the encoding operation for the semantic feature maps.

Specifically, for the remapping operation $\mathcal{R}$, we employ $3\times 3$ convolutional layers with a stride of 2 and padding of 1, each followed by a batch normalization layer and a rectified linear unit (ReLU) activation layer, adapting the feature map channels to 64, 256, and 512, respectively. Regarding the disparity encoding operation $\mathcal{E}^{D}$, we employ ResNet-152 [47] as the backbone network to extract features from the last disparity map $\boldsymbol{D}_{n}$. In ResNet-152, the first block consists of a convolutional layer, a batch normalization layer, and a ReLU activation layer. Then, a max pooling layer and four residual layers are sequentially applied to progressively increase the number of feature map channels.

Similarly, we utilize the residual block for the encoding operation $\mathcal{E}$ on the semantic feature maps, resulting in feature maps with 1024 and 2048 channels. The fused features $\mathcal{F}^{F}$ contain both texture and spatial geometric information, thereby enhancing semantic scene understanding. We conduct an ablation study for different feature fusion modules in Sect. IV-F.

We employ the decoder introduced in our previous work SNE-RoadSeg [3] to decode the fused features and generate the semantic prediction. In this encoder, three convolutional layers in the feature extractor and the upsampling layer share the same parameters: a $3\times 3$ kernel size, a stride of 1, and a padding of 1. In the final layer, features are upsampled to create the prediction map with $N$ channels, where $N$ denotes the number of semantic classes.

The loss function employed in our joint learning framework should guide the supervision of both the semantic segmentation and stereo matching tasks. Gradient smoothness between the disparity and semantic segmentation maps typically aligns closely, particularly at inter-class boundaries, where traditional training strategies tend to result in more errors due to factors such as occlusion and reflection. In light of this, we propose an SCG loss function to supervise the entire joint learning process, which leverages semantic consistency to optimize the training of S³M-Net.

Given the ground-truth semantic segmentation map $\boldsymbol{M}^{G}\in\mathbb{R}^{H\times W}$, Each pixel $\boldsymbol{p}$ of $\boldsymbol{M}^{G}$ can be written as follows:

where $C$ refers to the number of the semantic classes. We construct an extended 3D volume ${\boldsymbol{V}^{3D}}\in\mathbb{R}^{H\times W\times C}$ using the following expression:

where $c$ represents the $c$-th channel in the volume, and $\delta$ denotes the Kronecker Delta function [48]. As a result, each channel of the volume can be regarded as a binary segmentation map of the $c$-th class. To emphasize semantic consistency, We use an average pooling operation for each channel to obtain the inter-class volume $\boldsymbol{V}^{I}\in\mathbb{R}^{H\times W\times C}$:

where $\mathcal{P}$ denotes the average pooling operation. Furthermore, we apply a normalization operation:

to obtain a normalized volume $\boldsymbol{V}^{N}\in\mathbb{R}^{H\times W\times C}$. We then map $\boldsymbol{V}^{N}$ to a semantic consistency-guided weight map $\boldsymbol{W}\in\mathbb{R}^{H\times W}$ through:

The total loss function

consists of an SCG semantic segmentation loss $\mathcal{L}_{ss}$ and an SCG stereo matching loss $\mathcal{L}_{sm}$. $\mathcal{L}_{ss}$ is formulated as follows:

where $N$ denotes the pixel number, $C$ represents the class number, $\hat{y}_{c}(\boldsymbol{p})$ denotes the predicted probability of $\boldsymbol{p}$ belonging to class $c$, $y_{c}(\boldsymbol{p})$ represents the ground-truth label for $\boldsymbol{p}$ in class $c$, and $\alpha$ denotes the loss weight. Based on the ablation study detailed in Sect. IV-F, we set the value of $\alpha$ to 0.1. Moreover, $\mathcal{L}_{sm}$ is formulated as follows:

where $\boldsymbol{D}^{G}$ represents the ground-truth disparity map and $\boldsymbol{D}_{i}$ denotes the $i$-th disparity map in $\mathcal{D}$. $\alpha$ is set to 0.1 and $\gamma$ is set to 0.9.

Since the training of our network requires both semantic and disparity annotations, we employ two public datasets to evaluate its performance: the vKITTI2 [27] dataset (synthetic yet large-scale) and the KITTI 2015 [28] dataset (real-world yet modest-scale). Their details are as follows:

• •

  The vKITTI2 dataset contains virtual replicas of five sequences from the KITTI dataset. It provides 15 classes for the semantic segmentation tasks. Dense ground-truth disparity maps are acquired through depth rendering using a virtual engine. In our experiments, we randomly select 700 pairs of stereo images, along with their semantic and disparity annotations to evaluate the effectiveness and robustness of our proposed S³M-Net, where 500 pairs are used for model training and the remaining 200 pairs are used for model validation.

• •

  The KITTI 2015 dataset contains 400 pairs of stereo images captured in real-world driving scenarios, with 200 pairs containing ground truth and the other 200 pairs lacking ground truth. It provides 19 classes for the semantic segmentation tasks (in alignment with the Cityscapes [44] dataset). Sparse disparity ground truth is obtained using a Velodyne HDL-64E LiDAR. In our experiments, we allocate 70% of the dataset for training, while the remaining portion is used as the test set.

Our experiments are conducted on an NVIDIA RTX 3090 GPU. The batch size is set to 1. The maximum disparity search range is set to 192 pixels. All images are cropped to 1000 $\times$ 320 pixels before feeding into the network. We utilize the AdamW [55] optimizer for model training, setting the epsilon and weight decay parameters to $10^{-8}$ and $10^{-5}$, respectively. The initial learning rate is set to $2\times 10^{-4}$. Training lasts for 100K iterations on the vKITTI2 dataset and 20K iterations on the KITTI 2015 dataset. We employ traditional data augmentation techniques to enhance the robustness of the models.

Since our proposed S³M-Net simultaneously performs semantic segmentation and stereo matching, we evaluate the performance of both tasks in our experiments.

We utilize seven evaluation metrics to quantify the performance of semantic segmentation: (1) accuracy (Acc), (2) mean accuracy (mAcc), (3) mean intersection over union (mIoU), (4) frequency-weighted intersection over union (fwIoU) [56], (5) precision (Pre), (6) recall (Rec), and (7) F1-score (FSc).

Additionally, we use two evaluation metrics: (1) the average end-point error (EPE) and (2) the percentage of error pixels (PEP), setting the tolerance for the latter to 1.0 and 3.0 pixels, respectively, to quantify the performance of stereo matching.

The qualitative experimental results on the vKITTI2 and KITTI datasets are presented in Figs. 3 and 4, respectively, while the quantitative experimental results on the vKITTI2 and KITTI datasets are given in Tables I and II, respectively. These results suggest that S³M-Net outperforms other SoTA single-modal and feature-fusion networks (including CNN-based and Transformer-based methods) across all evaluation metrics on both datasets. Specifically, it is noteworthy that when the entire joint learning framework is trained by minimizing our proposed SCG loss, S³M-Net achieves the best performance on the KITTI dataset across all evaluation metrics except for Pre. Compared with SoTA methods, it shows improvements of $5.71\%$ in mAcc, $4.84\%$ in mIoU, $1.35\%$ in fwIoU, and $0.76\%$ in FSc, respectively. Similarly, it outperforms other networks on the vKITTI2 dataset in most evaluation metrics, with improvements of $1.72\%$ in fwIoU and $1.44\%$ in FSc. However, for mAcc, mIoU, and Rec, its performance is comparable to that of S³M-Net trained without using the SCG loss. Additionally, it is obvious that feature-fusion networks consistently outperform single-modal networks, particularly under challenging weather and lighting conditions. This observation aligns with our expectations, as feature-fusion networks leverage both RGB images and disparity maps, allowing them to effectively learn informative spatial geometric representations. However, SoTA feature-fusion networks may exhibit higher error rates in distant regions. For instance, FuseNet and SNE-RoadSeg demonstrate poor performance in the sky. We attribute this phenomenon to the deep structure of the encoders, where distinguishing distant objects using disparity features becomes challenging, and the feature fusion process amplifies the influence of the disparity feature. In contrast, within our proposed joint learning framework, we can extract more informative features benefiting from both tasks, irrespective of dataset size. This improvement is likely due to the fact that joint learning of multiple interconnected tasks introduces a form of regularization, which has shown its superiority over uniform complexity penalization in reducing over-fitting.

The qualitative experimental results on the vKITTI2 and KITTI datasets are given in Figs. 5 and 6, respectively, while the quantitative experimental results on the vKITTI2 and KITTI datasets are presented in Tables III and IV, respectively. These results suggest that S³M-Net outperforms other SoTA stereo matching networks across all evaluation metrics on both datasets. Specifically, S³M-Net trained with and without using the SCG loss achieves the top and second-best overall performances, respectively. S³M-Net, when trained without using the SCG loss, demonstrates improvements of $2.50\%$-$71.32\%$ in EPE, $4.17\%$-$64.19\%$ in PEP 1.0, and $4.49\%$-$67.97\%$ in PEP 3.0. On the other hand, S³M-Net, when trained with the SCG loss, shows improvements of $5.00\%$-$72.06\%$ in EPE, $5.44\%$-$64.38\%$ in PEP 1.0, and $4.49\%$-$69.83\%$ in PEP 3.0. We attribute these improvements to the feature sharing and fusion strategies applied in S³M-Net. First, sharing features with the semantic segmentation task allows S³M-Net to learn stereo matching effectively even with limited training data. Second, as discussed above, stereo matching can sometimes produce ambiguous disparity estimations, especially in occluded or texture-less areas. The pursuit of semantic consistency helps resolve such ambiguities, leading to more reliable disparity estimation results. In Fig. 6, it is evident that regions lacking disparity ground truth frequently have substantial errors. Previous stereo matching algorithms have endeavored to tackle this issue through knowledge distillation with pre-trained models [38]. Nevertheless, our S³M-Net successfully overcomes this challenge by leveraging semantic information.

In this subsection, we first conduct an ablation study on the selection of loss weight $\alpha$ in (11). Fig. 7 shows the quantitative experimental results with respect to different $\alpha$ in the range of $0.0$ to $0.4$ for both semantic segmentation and stereo matching. It can be obvious that when $\alpha=0.1$, S³M-Net achieves the best overall performance for both tasks. Further weight tuning is possible, but it should be approached cautiously, especially when dealing with limited data to avoid over-fitting.

Furthermore, we conduct an additional ablation study on the feature fusion strategy in our proposed FFA module. As shown in Table V, when using the addition operation to fuse heterogeneous features, the FFA module consistently achieves the best performance across all evaluation metrics, compared to other feature fusion strategies, including concatenation, cross feature module (CFM) [57], dynamic dilated pyramid module (DDPM) [58], separation-and-aggregation gate (SA Gate) [59], and softmax weighted sum (SWS) used in AysmFusion [60], CEN [61], and TokenFusion [62].