Lightweight inception-UNet with attention mechanisms for semantic segmentation

The encoding phase extracts features and reduces spatial dimensions. For this, the proposed LWA-MoDUNet employs only two encoding blocks for the encoding path. Each encoding block performs a comprehensive set of operations by incorporating multiple convolution layers, and max-pooling layers, with an inception block, and a batch normalization. Figure 2 illustrates the structure of an encoder block.

Figure 2:

Architecture of the encoder block of the proposed LWA-ModUNet.

From Figure 2, it is seen that the input data undergoes processing through three paths. The first path includes data processing through a 3 × 3 convolution layer and then through the inception module. The inception block consists of multiple kernels that capture features at different receptive fields. This ensures that the model has a better context than the standard convolution layer. The complete block of inception block is depicted in Figure 3 The features extracted from the inception module are passed to the decoding phase as well as within the same encoder block which enables the network to efficiently differentiate between the multiple-object boundaries and context, prevents information loss and mitigates the noise and occlusions in an image. In the second path, the 1 × 1 convolution layer, batch normalization, and 2 × 2 max-pooling process the input data. The processed data are included with extracted features from the inception module. This encourages the model to ensure that the learned features are more enriched than the original input in varied conditions for better scene understanding. Finally, the third path employs 2 × 2 max-pooling to downsample the input data and concatenate with the features from the second path. This acts as an input reinforcement to the next block and also aids the further layers to always have a comparative sense between the learned and the original features. This architecture of the encoder block enhances the network’s ability to discern crucial information from input images and distinguish between different objects in better scene understanding.

Figure 3:

Architecture of the inception block of the proposed LWA-ModUNet [54].

The decoding phase upsamples the feature space to produce a high-resolution segmentation mask. In the proposed architecture, only two decoder blocks are incorporated in the decoding phase. Each decoder block employs an attention module, a deconvolution layer, and convolution layers to perform the feature upsampling. The illustrative structure of the decoder block is depicted in Figure 4. In the decoder block, the features from the encoding phase are passed to the attention module to excite the important features at the channel level. The skip connections used in UNet provides spatial information from the decoder to the corresponding encoder but they bring along the poor feature representation. Hence, the modified inception has been integrated to enhance UNet with attention module as depicted in Figure 4 to reduce the complexity during feature extraction in an image.

Figure 4:

Decoder architecture of the proposed LWA-MoDUnet.

The attention block [55], as depicted in Figure 5, focuses on attention coefficients β_i ∈ [0, 1], which identifies the significant features and then assign the weights to them. The attention coefficient is element wise multiplied with generated input feature maps, that is, mil m_i^l , c=mil c = m_i^l , c ⋅=βil c\; \cdot = \beta _i^l , to produce the result from attention gates. In a default configuration, the value of single scalar attention is determined for each pixel vector mil∈RNl m_i^l \in {R^{Nl}} in L layer for N number of feature maps. The gating vector v_i ∈ R^Nv that contains the contextual information significantly reduces the influence of less relevant features. The gating coefficients are obtained by using additive attention, which has been equated in Eqs (1) and (2): (1) qattl=ΨT σ1WT mil+WT vi+bv+bψ q_{att}^l = {\Psi ^T}\;{\sigma _1}\left( {{W^T}\;m_i^l + {W^T}\;{v_i} + {b_v}} \right) + {b_\psi } (2) βl,i=σ2qattlmil,vi;Θatt {\beta ^l},i = {\sigma _2}\left( {q_{{\rm{att}}}^l\left( {m_i^l,{v_i};{\Theta _{{\rm{att}}}}} \right)} \right) where σ₂(m_i,c) denotes the sigmoid activation, Θ_att depicts the parameters used to characterize attention gates that include W_m ∈ R^Nl × Nint and W_v ∈ R^Nl×Nint. For input tenors, a channel-wise convolution of 1 × 1 × 1 has been performed to calculate the linear transformations. b_ψ ∈ R and b_v ∈ R^Nint denote the bias used.

Figure 5:

Diagram of additive attention gate [55].

The sequential use of softmax activation [56], used in attention block, is to normalize the attention coefficients. Due to the implementation of sigmoid activation, the proposed model has shown better convergence during training. Within the architecture, skip connections are utilized to emphasize important features. The updating rule for convolution parameters from the lower layers (l − 1) is shown in Eq. (3): (3) ∂m^l∂βlfml−1;ϕl−1∂fml−1;ϕl−1∂βli=i=βl+i ml∂ϕl−1∂ϕl−1i ∂ϕl−1∂ϕl−1i \matrix{ {\partial \left( {{{\hat m}^l}} \right)} \hfill \cr {\partial \left( {{\beta ^l}f\left( {{m^{l - 1}};{\phi ^{l - 1}}} \right)} \right)} \hfill \cr {\partial \left( {f\left( {{m^{l - 1}};{\phi ^{l - 1}}} \right)} \right)} \hfill \cr {\partial \left( {{\beta ^l}} \right)} \hfill \cr {i = i = {\beta ^l} + i\;{m^l}} \hfill \cr {\partial \left( {{\phi ^{l - 1}}} \right)} \hfill \cr {\partial \left( {{\phi ^{l - 1}}} \right)} \hfill \cr {i\;\partial \left( {{\phi ^{l - 1}}} \right)} \hfill \cr {\partial {{\left( {{\phi ^{l - 1}}} \right)}^i}} \hfill \cr }

This focuses on capturing the relevant contextual information between the regions of an image. Furthermore, the attention module can effectively fuse features from multiple scales within an image. This benefits in capturing both local and global contexts, making it suitable for objects of varying sizes, which is particularly valuable for distinguishing objects from their surroundings and handling complex scenes. Consecutively, a couple of standard convolution layers are followed by a deconvolution layer in a decoder block. This advantage in suppressing the noise in the predicted segmentation mask and pixelwise prediction is for accurate segmentation of complex images. Finally, the decoder block has a 1 × 1 convolution layer that processes the features from the skip connections and performs concatenation with the upsampled features by the deconvolution layer. This helps the model to retain the coarse features from the encoder side. At the end of the decoding phase, a 1 × 1 convolution layer is applied to generate the output segmentation mask. The flowchart of the proposed model is shown in Figure 6, depicting the generation of the segmentation mask.

Figure 6:

Flowchart of the proposed LWA-MoDUNet depicting the encoding and decoding phases with inception and attention modules.

To assess the performance of the proposed semantic segmentation method, five publicly accessible datasets, namely, Autorickshaw Detection Dataset [18], IDD-Lite [19], CT-Liver [20], MADS [21], and PETS [22], are used. The used datasets contain objects within diverse images of scenes characterized by congested and unstructured patterns, which make them ideal for the assessment. The Autorickshaw Detection Dataset [18] consists of 1,000 images and has been publicly released by the Indian Institute of Information Technology, Hyderabad, India, in conjunction with the Autorickshaw Detection Challenge. The IDD-Lite [19] comprises 1,404, 204, and 408 training, validation, and testing images, respectively. This dataset is organized into seven distinct classes, namely, drivable regions, non-drivable areas, living entities, vehicles, roadside objects, distant objects, and the sky. The PETS [22] dataset consists of 37 sets of various breeds of cats and dogs in varied lighting, poses, and background conditions. Each set of breeds contains approximately 200 images, resulting in a total of 7,349 images that have been divided into training, testing, and validation datasets. From 7,349 images, 3,312 have been taken for training, 368 for validation, and 3,669 for testing purposes. The CT-Liver [20] dataset consists of 232 images in which 152 images have been considered for training, 35 for testing, and 45 for validation purposes. The MADS [21] dataset is a collection of human movement images of martial arts, dance, and sports movements in varied angles for pose detection. For training, validation, and testing of segmentation, 861, 152, and 179, respectively, images have been taken. Figure 7 depicts the representative images from each dataset along with the respective ground-truth (GT) images.

Figure 7:

Sample images of considered datasets (Row 1) along with their corresponding GT (Row 2). GT, ground truth; IDD-Lite, Indian driving dataset-lite; MADS, martial arts, dancing and sports.

The quality of the generated segmentation masks for the abovementioned datasets has been comprehensively evaluated against seven performance parameters, namely, IoU score, error (Err), specificity (Sp), sensitivity (Ss), accuracy (Acc), correlation coefficient (Cc), and F-score.

During training, each image is first resized to p × q, where p = 256, q = 256, and randomly flipped horizontally, and croppped m × n, where m = 128 and n = 128. The input images were subjected to different augmentations to avoid distortions and over-fitting such as cropping, zoom-in/out, rotation, lightning, and affine (scaling and translation). For the data normalization, each image is subtracted from the mean and divided by standard deviation. Each input image is resized to 128 × 128 × 3 dimensions, where 3 represents the RGB channel. For the autorickshaw, PETS, MADS, and CT-Liver datasets, the loss function employed is the flattened binary cross-entropy loss [57], also known as flattened sigmoid cross-entropy loss. As the IDD-Lite is a multi-class dataset, cross-entropy loss [58] is utilized. In LWA-MoDUNet, since both the instances of feature and target networks are convolutional layers, inspired from the idea mentioned above to rectify the gradient vanishing issue, we have used Xavier initialization [59] for convolutional layers to speed up the training process. Surprisingly, no existing backbone is used as all of the model is trained from scratch by training over 200 epochs. The Adam optimizer [60] with early stopping and step learning rate scheduler with initial learning rate as 1E − 3 is used. All other hyperparameters are considered as default. At the output, bilinear interpolation is employed to resize the predicted saliency maps. The model is implemented in PyTorch 2.0.0 and trained on an Intel Xeon 2.2 GHz CPU (24GB RAM) and an NVIDIA A100-SXM4 (40 GB memory), which are a Hexa-core, 12-thread PC.

To explore the robustness of the proposed model, an ablation study [61] is conducted through a quantitative analysis of the proposed model with and without inception and attention modules. The hyperparameter settings of the proposed model are detailed in Section “Implementation details”. Experimental evaluation is performed on the autorickshaw dataset by considering four metrics, that is, IoU score, Error, Accuracy, and F-score. The experimental results are depicted in Table 1.

It can be observed from Table 1 that the inclusion of the inception module has significantly enhanced the performance of the LWA-MoDUNet (proposed model) as compared with its absence. Furthermore, the model has attained superior performance with the use of the attention module as shown in Table 1. Therefore, this study supports the architectural advantages of the proposed model and is evident from the efficient results.

In this section, detailed analysis of the results obtained are presented. Examining Figure 8, it becomes apparent that the comparative SOTA, especially UNet with ResNet-18 encoder and E-Net, struggled in the segmentation of the input images, exhibiting noise, poor-boundary definitions, and gaps. While LWA-MoDUNet can generate similar segmentation masks that best match the GT. The resultant segmentation masks have sharp and smooth boundaries with minimal noise. In Figure 9, noticeable discrepancies were observed in the segmented results of E-Net and Seg-Net. Conversely, LWA-MoDUNet consistently generated outputs closely resembling their respective GT images for the examined figures. Furthermore, the output images of the LWA-MoDU-Net demonstrated comparatively lower levels of distortion. Similarly, from Figure 10, it can be observed that the results obtained from the LWA-MoDUNet are approximately similar to the GT of the respective images under consideration. The segmentation masks have accurate boundaries even in the cluttered backgrounds, while ResNet18 and ENet models generated segmentation masks that were partially matching with the GT of the original image. Conversely, UNet and SegNet underperformed by generating distorted output masks. On the contrary, considering Figure 11, it can be witnessed that LWA-ModUNet and UNet with ResNet34 have achieved competitive results. The segmented outputs of the earlier one are better than the later one. While UNet underperformed and outputs the distorted masks. Moreover, Figure 12 shows that the considered models have shown average results. The result suggests that the masks produced by LWA-ModUNet are considerably similar to their original GT images. The compared models have generated distorted images and hence cannot segment the image effectively.

Table 2 presents the experimental results on the autorickshaw dataset. UNet with ResNet18/ResNet34 encoder demonstrate similar performance across all the parameters. However, their performance falls behind that of the other models under consideration, as well as the proposed methods for the autorickshaw dataset. Notably, LWA-MoDUNet outperformed comparative models, achieving the highest IoU-score, that is, 0.8768. Furthermore, LWA-MoDUNet demonstrated superiority across all considered parameters, that is, error as 0.0663, accuracy as 0.9336, specificity as 0.9429, sensitivity as 0.9248, F-score as 0.9344, and correlation coefficient as 0.8676. Similarly, upon analyzing the results for the IDD Lite dataset in Table 3, it is evident that our proposed model exhibited exceptional performance, surpassing the state-of-the-art comparative models, achieving the highest IoU score of 0.9283. The second highest IoU of 0.6174 is achieved by UNet-ResNet34, which shows that the proposed model segments the complex images efficiently.

Similarly, the performance of the compared models and proposed model on the PETS dataset has been illustrated in Table 4. The proposed LWA-ModUNet performed better than the compared SOTA models, indicating the highest IoU score of 0.8768. On other parameters, LWA-ModUNet performed commendably achieving the highest values for accuracy, sensitivity, specificity, F-score, and correlation-coefficiency. According to the result, LWA-MoDUNet is most suitable for accurately segmenting objects in complex images, achieving and maintaining a better balance between precision and recall. Contrarily, SegNet consistently delivers lower performances on all parameters across the metrics, thus making the model a crucial choice for the semantic segmentation of complex images. By contrast, UNet with ResNet34 and UNet with ReSNet18 performed better but remained behind LWA-MoDUNet reporting an IoU score of 0.8459 and 0.8356, respectively.

The performance evaluation of the models on the MADS dataset is illustrated in Table 5. The results depict that the performance of SegNet is comparable to the LWA-MoDUNet in terms of all the considered parameters. In terms of IoU score, LWA-MoDUNet reported 0.9768, which is the highest while SegNet achieved 0.9721. UNet underperformed by achieving the lowest IoU score value of 0.8968. Table 6 illustrates the segmentation results of the compared models on the CT-Liver dataset. The table suggests that the proposed model marked the highest parametric values across the metrics, achieving an IoU score of 0.9651, accuracy of 0.9821, correlation coefficient of 0.9643, specificity of 0.9855, F-score of 0.9822, and sensitivity of 0.9788, respectively, which are the highest. UNet with ResNet34 and UNet with ResNet-18 fall behind LWA-MoDUNet with an IoU score of 0.9445 and 0.9565, respectively. On the contrary, SegNet, and Enet fall slightly behind the ResNet encoders resulting in IoU score of 0.9495 and 0.9349, respectively. Moreover, UNet underperformed in the segmentation of image reporting the lowest IoU score of 0.8839. Hence, based on the segmentation accuracy (IoU score), it can be confidently affirmed that our proposed model represents a substantial improvement, attributable to its enhanced capacity to extract deeper features.

Figures 13–17 demonstrate that LWA-MoDUNet attains the best IoU score, exhibiting a substantial margin of superiority over the compared models for all five compared datasets. Furthermore, for visual illustrations, the graphs for the models’ performance at the validation phase provide unequivocal evidence that the proposed method consistently outperforms throughout the entire span of epochs considered. It is worth noting that the proposed model maintains its supremacy in terms of IoU score and outperforms the compared state-of-the-art-methods.

Figure 13:

Comparative analysis of IoU score on training and validation images taken from autorickshaw dataset. IoU, intersection over union.

Figure 14:

Comparative analysis of IoU score on training and validation images taken from IDD-Lite dataset. IDD-Lite, Indian driving dataset-lite; IoU, intersection over union.

Figure 15:

Comparative analysis of IoU score on training and validation images taken from PETS dataset. IoU, intersection over union.

Figure 16:

Comparative analysis of IoU score on training and validation images taken from MADS dataset. IoU, intersection over union; MADS, martial arts, dancing and sports.

Figure 17:

Comparative analysis of IoU score on training and validation images taken from CT-Liver dataset. IoU, intersection over union.

To validate, the same Wilcoxon rank sum test has been conducted by using the IoU scores of the two models across varied sample sizes, that is, 30 and 50. If the p-value is lower than the significance level (α = 0.05), then the null hypothesis (i.e., both models are statistically similar) is rejected, and the alternative hypothesis (i.e., the proposed model is statistically superior to the compared model) is accepted. Table 7 depicts the p-values of the proposed model in comparison to other considered models. Notably, it is clearly evident from the Table 7 that the p-value of the compared models is considerably less than the α. Furthermore, it can be observed in Table 7 that the p-value is indirectly proportional to the sample size. Therefore, it is pertinent to state that the proposed model performs statistically better segmentation in comparison to the compared models.

Additionally, the results in Figure 18 clearly show that LWA-ModUNet yields fewer Type 1 errors than the compared state-of-the-art models. On the contrary, the proposed model depicted comparable Type 2 errors on PETS and IDD datasets but lower errors on the remaining datasets as compared with the state-of-the-art models. As compared -with the UNet and ResNet-based models LWA-MoDUNet raises less false alarms.

Figure 18:

Model type 1 and type 2 errors comparison across datasets. IDD-Lite, Indian driving dataset-lite; MADS, martial arts, dancing and sports.

Furthermore, this study analyzed the complexity of these models in terms of parameters and inference time. As illustrated in Table 8, LWA-ModUNet has significantly fewer parameters among the considered models. On the contrary, DeconvNet exhibits the highest number of parameters, along with the training and inference times. The inference time for ResUNet is longer than that for LWA-ModUNet, indicating the drawbacks of deep residual learning. Furthermore, U-Net has nearly six times more parameters than LWA-ModUNet despite being based on the UNet architecture. When compared with FCN-8s and SegNet, the inference time of LWA-ModUNet is comparatively less than both. The simpler structure of LWA-ModUNet has contributed to reduced execution time. With improved accuracy and minimal computational resources, LWA-ModUNet demonstrates superior performance. Hence, LWA-ModUNet is a computationally lightweight model, making it desirable for real-time applications.