Deep Learning for Computer Vision: My Journey into Semantic Segmentation

Deep learning has revolutionized computer vision, unlocking possibilities in fields ranging from autonomous driving to medical imaging. Recently, I had the incredible opportunity to dive deep into this transformative domain through the Goal-Oriented Free Mentorship program with an industry expert, **Dr. Olabode Sule**, who has worked at top-tier tech companies like Meta and several startups. This experience allowed me to not only deepen my understanding of semantic segmentation but also to refine essential skills in implementing state-of the-art neural network architectures and enhancing model performance.

Project Overview

Semantic segmentation, a key task in computer vision, involves assigning a label to every pixel in an image. It’s widely applied in real-world scenarios such as road scene understanding, where models need to distinguish between roads, vehicles, pedestrians, and other elements. Beyond autonomous driving, semantic segmentation has transformative applications in healthcare, such as identifying tumors in medical imaging or segmenting organs for surgical planning. In the energy sector, it is used to detect defects in manufacturing processes like batteries production, enabling predictive maintenance and quality control.

My project focused on designing and training semantic segmentation models using popular architectures like U-Net¹ and Deeplabv3². Using the Mapillary dataset, a diverse collection of street-level imagery, I trained and evaluated models to achieve high accuracy in segmenting complex scenes. The dataset’s diversity posed challenges, but it also offered an excellent opportunity to experiment with techniques that could generalize well across varied environments.

Additionally, I gained hands-on experience with the Google Cloud Platform (GCP), including setting up google cloud virtual machines with GPU instances to accelerate training. I used Google Cloud Storage (GCS) buckets for efficient data management, and wrote scripts in Google Cloud Command Line Interface (gcloud CLI) to streamline operations and organize workflows. These tools proved invaluable in scaling the project and ensuring smooth experimentation.

In the following, I will give an overview of the concepts and architectures that I learned while working on this project.

Understanding Fully Convolutional Networks

Fully convolutional networks (FCNs)³ are neural network architectures composed entirely of convolutional operations, often accompanied by activation functions, batch normalization, and dropout layers. This design makes them highly effective for tasks requiring spatial information, such as semantic segmentation. Unlike traditional convolutional neural networks (CNNs)⁴, which often include fully connected layers for classification, FCNs replace these with convolutional layers, ensuring that the output remains spatially aligned with the input image, preserving its dimensions.

To simplify: • Traditional CNNs: Often reduce the image to a fixed-size vector, losing spatial alignment. For instance, they might take a 256 x 256 image and output a single vector representing class probabilities. • Fully Convolutional Networks: Process the image entirely through convolutional operations, enabling pixel-level outputs where the dimensions of the input and output images are preserved. This is critical for tasks like semantic segmentation, where the label for each pixel matters.

Imagine trying to paint over a photo. In a fully convolutional approach, every pixel is painted individually but retains its position in the image. This spatial consistency is a defining feature of FCNs. See a schematic of an FCN depicted in Figure 1.

Figure 1. Schematic of a Fully Convolutional Network (FCN) for semantic segmentation. The input (left) is an image containing a cat and a dog. The intermediate blocks represent convolutional layers and nonlinear activations that extract hierarchical features. The final output (right) is a pixelwise prediction map, where each pixel is classified into a semantic category. FCNs use upsampling techniques to restore spatial resolution and produce segmentation maps that are aligned with the original image. This figure is from Ref. [3].

U-Net Architecture

The schematic of this architecture is best represented in Figure 2, depicting a U-shaped diagram which can be categorized into three main components:

1. Contracting Path (Left Side of the U): Like zooming into details, this path identifies features at smaller scales by downscaling the image through convolution and pooling operations. Let’s break down convolution and pooling in simple terms: • Convolution: This operation, depicted as blue arrows in Figure 2, scans small portions (patches) of the image using filters to extract meaningful patterns such as edges or textures. Imagine looking at a photo through a magnifying glass that highlights specific features (e.g., a car’s edges). • Pooling: This reduces the size of the image while keeping the most important information. Think of it as summarizing a large page of text into a few key sentences. In U-Net, pooling helps focus on high-level patterns by progressively reducing the image size. This is shown as red arrows.
2. Expanding Path (Right Side of the U): This path reconstructs the image by progressively upsampling (green arrows in the figure), restoring spatial details. It ensures that the output has the same dimensions as the input while improving localization accuracy.
3. Skip Connections: A crucial feature of U-Net that directly links corresponding layers from the contracting path to the expanding path, as shown as grey arrows in Figure 2. Think of it as copying details from the left side of the U and pasting them onto the right side, preserving clarity and precision. These connections have been empirically determined to help with details such as delineating the boundaries of images.

Figure 2. U-net architecture (example for 32x32 pixels in the lowest resolution). Each blue box corresponds to a multi-channel feature map. The number of channels is denoted on top of the box. The x-y-size is provided at the lower left edge of the box. White boxes represent copied feature maps through the skip connection protocol. The arrows denote the different operations. Figure adapted from Ref. [1]

Deeplabv3 Architecture

In some fully convolutional network (FCN) architectures, repeated max-pooling and striding across layers significantly reduce the spatial resolution of feature maps. To overcome this limitation and generate denser feature maps, the final downsampling operations are removed, and instead, the filters in subsequent convolutional layers are upsampled. This technique, known as atrous convolution, allows convolutions with expanded receptive fields without increasing the number of parameters or computational cost.

Compared to standard convolutions with large kernels, atrous convolutions enlarge the field of view more efficiently. Furthermore, to address the challenge of detecting objects at multiple scales, multiple parallel atrous convolutions with different dilation rates are applied — a method known as Atrous Spatial Pyramid Pooling.

In other words, the architecture of Deeplabv3 as illustrated in Figure 3 include two main elements:

1. Atrous (Dilated) Convolutions: Atrous convolutions are like looking through a mesh sieve with adjustable gaps. By expanding the field of view without increasing computation, these convolutions help the model capture patterns at varying scales — such as distant trees and nearby pedestrians in the same image. They are crucial for preserving spatial resolution in feature maps, making them ideal for segmenting large, complex scenes.
2. Atrous Spatial Pyramid Pooling (ASPP): ASPP is an advanced technique that applies multiple atrous convolutions with different rates in parallel. Imagine using sieves of various sizes simultaneously to extract features at small, medium, and large scales. This approach ensures that the model captures both fine details and broader context, which is vital for segmenting objects of different sizes in an image.

Figure 3: The architecture used in DeepLabV3. Starting from an input image, a backbone network (e.g., ResNet) is used to extract feature maps. This backbone progressively reduces spatial resolution through strided convolutions, eventually reaching an output stride of 16 (shown through Blocks 1–4). At the final feature map (Block4), the ASPP module is applied. This consists of: (a) Four parallel atrous convolutions with different atrous rates (b) Image-level pooling, which captures global context. All these parallel outputs are concatenated, followed by a 1×1 convolution to fuse features. This image is gotten from Ref. [2]

Next, I would like to share a list of the skills that I acquired and the tools utilized while working on this project.

Key Achievements of this Project

1. Exploratory Data Analysis (EDA) and Visualizations.

• Conducted an in-depth EDA to understand the dataset’s structure and characteristics.
• Used Plotly and Matplotlib to visualize class distributions, basic visual statistics (e.g., pixel intensity histograms, class frequency plots), and image samples.

2. Model Exploration and Implementation:

• Implemented fully convolutional architectures, including U-Net and Deeplabv3, by adapting open-source repositories using PyTorch and CUDA.
• Gained hands-on experience in adapting these architectures for training on the Mapillary dataset.

3. Enhancing Model Performance:

• Validation and testing scores were calculated based on the Multiclass Dice coefficient, which is the average Dice coefficient across all classes. The Dice coefficient measures the similarity between predicted and ground truth regions, effectively evaluating how well the model segments different regions. Here is a good resource on understanding dice coefficient⁵.
• Designed custom data augmentation pipelines using the Albumentations⁶ library, including transformations like random cropping, flipping, and color adjustments, to improve model generalization.
• For the loss function, utilized a combo loss function: cross-entropy and dice loss⁷. While cross-entropy focuses on accurately predicting pixel-level classifications, dice loss specializes in accurately segmenting regions in the image. This combination balanced both pixel-level detail and region-level accuracy.

4. Optimization Technique:

• Employed the RMSprop optimizer, a gradient-based optimization algorithm designed to maintain a moving average of squared gradients and scale the learning rate accordingly. This approach effectively adjusts the step size during training to optimize the combo loss function.
• By addressing the challenge of oscillating gradients, RMSprop helps stabilize and accelerate convergence. For a detailed tutorial, refer to this resource on RMSprop optimizer.

5. Experiment Tracking and Reporting:

• Leveraged Weights and Biases (wandb) to track experiments, compare results, and analyze metrics comprehensively.
• Utilized wandb’s visualization tools to identify trends, debug issues, and communicate insights effectively.

6. Cloud Computing and Workflow Optimization:

• Set up Google Cloud virtual machines with GPU support to expedite training⁸.
• Utilized GCS buckets for seamless data handling and retrieval.
• Developed scripts in gcloud CLI to automate operations, organize resources, and improve project efficiency.

7. Code Quality and Modularity:

• Developed modular and maintainable Python code, emphasizing reusability and clarity.
• Implemented robust training pipelines to streamline experimentation.

Now, let’s talk about the main results

Results

The following results focus on the training and performance of the Deeplabv3 semantic segmentation model over 20 epochs, showcasing visual evidence from segmentation outputs with quantitative trends in train loss and validation score.

Segmentation Outputs Across Epochs: Figure 4 shows the segmentation outputs of the model’s ability to learn and predict accurate masks for road scenes. These masks represent the actual segmentation labels, where each gray intensity or color corresponds to a specific class (e.g., road, cars, buildings, vegetation). Across epochs, we observe the following:

Figure 4: Five representative epochs of images, ground truths, and model predictions

• At epoch 6: The predicted masks are blurry and lack clear boundaries, indicating the model is in its initial learning phase and has yet to understand object and class distinctions.
• At mid-epochs (epoch 13 to 16): Predictions improve significantly as the model begins to recognize larger structures (e.g., roads, vehicles, and buildings). Boundaries are sharper, and details become more refined.
• At epoch 20: The predicted masks align closely with ground truth masks, showcasing the model’s ability to generalize to complex scenes. This reflects a high level of performance, with detailed and somewhat accurate segmentations.

Training Loss and Validation Score Trends: The quantitative metrics that provide a deeper understanding of the model learning process are presented in Figure 5.

Figure 5: Train loss and validation score trends over 20 epochs. Key points include epoch 6: 0.0718 (0.7428), epoch 13: 0.0572 (0.7484), epoch 15: 0.0561 (0.7459), epoch 16: 0.0556 (0.7487), and epoch 20: 0.0543 (0.7457) for train loss (validation score), respectively.

• Train Loss: A steadily decreasing loss over epochs suggests that the model effectively minimizes errors on the training set. The sharp decline in early epochs reflects rapid feature learning, while the plateau after epoch 18 indicates convergence.
• Validation Score: The validation score shows fluctuations in early epochs (common during initial feature extraction), followed by a steady improvement as the model generalizes better. The score peaks around epochs 15–18, aligning with the qualitative improvement in predicted masks seen in the segmentation figure. The slight decline at epoch 20 suggests potential overfitting, where the model might prioritize the training set at the expense of unseen data.

Recommendations for Optimization

1. Hyperparameter Tuning: Adjust learning rate and batch size to refine the model’s convergence and validation stability. Also, experiment with various loss functions to address class imbalances and improve pixel-level accuracy.
2. Pre-training: Increasing the number of training epochs does not necessarily improve model performance. In fact, as seen in Figure 5, the validation score starts to decline toward the end of training, suggesting potential overfitting. Instead of simply training longer, one alternative is to pretrain some model layers on a related task, such as ImageNet (or another large dataset). Additionally, modern self-supervised pre-training techniques can provide a strong initialization, improving generalization. Some of these techniques include, MoCo⁹, SimCLR¹⁰ or the more powerful DinoV2¹¹.

Note that the metric trends for the U-Net model (results not shown) are similar to the Deeplabv3 model. For the U-Net, the best epoch was 17 with a validation score of 0.7583, whereas the Deeplabv3 model was optimal at epoch 18 with a validation score of 0.7512.

Summary

This project was an eye-opener in several ways. Working on semantic segmentation helped me appreciate the nuances of neural network design and training. I learned how critical it is to choose appropriate loss functions and data augmentation techniques for specific tasks. Experiment tracking with wandb became an indispensable part of my workflow, allowing me to stay organized and iterate efficiently.

Moreover, the mentorship experience highlighted the importance of writing clean, modular code. Building maintainable projects not only makes collaboration easier but also sets a strong foundation for scaling solutions. The experience with GCP further enhanced my ability to leverage cloud platforms for scalable and efficient deep learning workflows.

This project has inspired me to further explore the intersection of deep learning and computer vision. I’m particularly interested in deploying these models for real-time applications and experimenting with transformer-based architectures such as Swin Transformer and Mask2Former. I also aim to explore semi-supervised and unsupervised techniques to reduce the reliance on large labeled datasets.

I am deeply grateful for the mentorship and guidance I received during this journey. It was an invaluable experience having the chance to learn from Dr. Olabode Sule, an expert in the fields of machine learning, deep learning, and artificial intelligence (AI). This has fueled my passion for innovation and solidified my commitment to mastering AI technologies.

Semantic segmentation is just one piece of the puzzle in deep learning, and I am excited to continue pushing boundaries in this ever-evolving field of deep learning and AI. If you are working on similar projects or have insights to share, I would love to connect and learn from your experiences! Feel free to check out the project on my GitHub or connect with me on LinkedIn to collaborate further!

References

1. Ronneberger, O., Fischer, P., & Brox, T. (2015). U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assisted intervention–MICCAI 2015: 18th international conference, Munich, Germany, October 5–9, 2015, proceedings, part III 18 (pp. 234–241). Springer international publishing.
2. Chen, L. C., Papandreou, G., Schroff, F., & Adam, H. (2017). Rethinking atrous convolution for semantic image segmentation. arXiv preprint arXiv:1706.05587.
3. Long, J., Shelhamer, E., & Darrell, T. (2015). Fully convolutional networks for semantic segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 3431–3440).
4. O’shea, K., & Nash, R. (2015). An introduction to convolutional neural networks. arXiv preprint arXiv:1511.08458.
5. For an interesting article in implementing the Dice Coefficient, see Understanding DICE COEFFICIENT.
6. Buslaev, A., Iglovikov, V. I., Khvedchenya, E., Parinov, A., Druzhinin, M., & Kalinin, A. A. (2020). Albumentations: fast and flexible image augmentations. Information, 11(2), 125.
7. Azad, R., Heidary, M., Yilmaz, K., Hüttemann, M., Karimijafarbigloo, S., Wu, Y., … & Merhof, D. (2023). Loss functions in the era of semantic segmentation: A survey and outlook. arXiv preprint arXiv:2312.05391.
8. To set up environments with required GPU hours for training and evaluating models, refer to this Reddit thread on how to obtain free trial GCP usage.
9. He, K., Fan, H., Wu, Y., Xie, S., & Girshick, R. (2020). Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 9729–9738).
10. Chen, T., Kornblith, S., Norouzi, M., & Hinton, G. (2020). A simple framework for contrastive learning of visual representations. In International conference on machine learning (pp. 1597–1607). PmLR.
11. Oquab, M., Darcet, T., Moutakanni, T., Vo, H., Szafraniec, M., Khalidov, V., … & Bojanowski, P. (2023). Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07193.