DEV Community

Cover image for Image Segmentation - Part 1: Data Exploration and Baseline Model (S1E1)
Levente Slajcho
Levente Slajcho

Posted on • Edited on

Image Segmentation - Part 1: Data Exploration and Baseline Model (S1E1)

#ai #pytorch #computervision #python

Table of contents

Foreword


Don't worry if you have never created a single model before.

This post is written for beginners with some theoretical knowledge, aiming to explain concepts in the practical way I wish I had found when I started going deeper with neural networks. That being said: your feedback means the world to me - especially if you are an expert in this field, reading this.

As mentioned in the Intro, we start our series with a Kaggle competition, solving an image segmentation problem, labeling cars on a Birds-Eye View (BEV) LiDAR map.

This project consists of 3 parts. By the end of this one, Part 1, you will build an intuition about:

  • U-Net Structure
  • Encoders and Decoders
  • Skip connections
  • IoU metric for evaluating segmentation masks
  • Understand what's COCO-style dataset

*This is intentionally a simple baseline model. The goal of Part 1 is not maximum accuracy, but understanding the full pipeline.
*
Example inference

Prerequisites

Theoretical:

Technical:

import torch
print(torch.cuda.is_available())
Enter fullscreen mode Exit fullscreen mode

If False, run this: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

GitHub Repository

🤖 Use this repository for this project

Make sure you're on the part-01-baseline-model branch.

For the best learning experience, I kindly ask you to copy the code by typing, preferably even without ML full-line autocomplete extensions. There are no shortcuts for understanding concepts and building intuition, your brain will thank you a lot for actively thinking while typing.

Let's Dive In

  1. Project Structure
  2. Loading the Dataset
  3. Building the Model
  4. Training the Model
  5. Exploration

Project Structure


In the first step, create the following project structure:

segment-vehicle-lidar/
|
├── data/
├── notebooks/
│   ├── 01_exploration.ipynb
|   └── 02_training_baseline.ipnyb
├── src/
│   ├── dataset.py
│   ├── model.py
│   └── train.py
├── checkpoints/
└── README.md
Enter fullscreen mode Exit fullscreen mode

Next:
🤖 Download the dataset from Kaggle

After extracting the file copy the test/ and train/ folders into the data/ folder.

Loading the Dataset

Open dataset.py.

import os #handles paths
import json #handles train.json
import numpy as np 
import cv2 #for loading and converting images and for handling polygons
import torch #does neural network stuff
from torch.utils.data import Dataset #handles data for neural network stuff
Enter fullscreen mode Exit fullscreen mode

Let's take a look at the code from the repository. In this file, you can see a class named CocoSegDataset(Dataset). This class tells us how and from where the COCO-style dataset should be loaded. Luckily, we don't have to implement the whole COCO format ourselves. For this project, our custom CocoSegDataset class reads the parts we need and turns the segmentation polygons into masks.

The COCO-style format is a standardized JSON-based format for training models for diverse tasks where we don't only want to know what we see in the image, but more importantly where we see it. A few examples are 2D human pose estimation, image labeling, and object detection.

Open the train.json annotation file. First you see a list of the metadata of the images and an assigned id. Then, when you scroll down, among some other information, you'll find the image_id, segmentation, bbox and area.

As you can see in the image of whatever I managed to draw below (let's call it Myrtle the Cat), segmentation (black) denotes a polygon in x1, y1, x2, y2, ... format. The area of the labelled data and bbox (blue) that contains [x, y, width, height] for the bounding box.

In the for loop of CocoSegDataset class's __getitem__ method, you can see how OpenCV is capable of turning these polygons into a 2D matrix mask.


COCO-style data annotation consisting of Myrtle the Cat and a polygon, bounding box and area

After reading the images, the last important step happens in these lines:

image = torch.from_numpy(image).permute(2, 0, 1).float()
mask = torch.from_numpy(mask).float().unsqueeze(0).float()
Enter fullscreen mode Exit fullscreen mode

The first line's purpose is to change the order of dimensions in a way PyTorch likes it. It expects the input channels (RGB) first, then the height and width of the image.

The second operation simply turns a 2D mask into a 3D array by using .unsqueeze(0) where 0 specifies the position of the newly added dimension. This dimension is needed to make mask have a channel dimension and match the number of dimensions of image.

Intuition Hook

In the beginning, after reading tons of sophisticated and very mathematic description of tensors, I realized that in Computer Science, a tensor is just a one-dimensional or multidimensional array. Nothing more, nothing less. Frameworks like PyTorch and TensorFlow mainly work with tensors.
Flat Jerry is a matrix, 3D Jerry is a tensor PyTorch likes.

  • .squeeze(d) flattens Jerry by removing its dimension d
  • .unsqueeze(p) inflates Jerry by adding a new dimension to the array at position p

Building the Model

Time to build our model!
We are going to implement this network,

Here, the Encoder - consisting of 3 EncloderBlocks - extracts the features in the image, and the Decoder - consisting of 3 DecoderBlocks - upsamples the intermediate features and creates the output - the yellow label over Myrtle the Cat.

Open model.py and copy the code into your project. You will see 4 classes defined, ordered by level of abstraction in the models structure:

  • DoubleConv(nn.Module): Implements a Module, which is the base class for all neural network modules, a fundamental building block for neural networks. For our baseline model in Part 1, we start building our model with Conv2d layers, each followed by a ReLU activation function.

  • The Encoder and the Decoder:

    • The EncoderBlock (marked red) consists of one DoubleConv module for feature extraction and a MaxPool2d for downsampling. In our project, these blocks learn increasingly abstract features that help the model identify which pixels belong to cars.
    • The DecoderBlock (marked blue) upsamples the feature maps using transposed 2D convolution, ConvTranspose2d (often incorrectly called deconvolution), followed by a previously defined DoubleConv module.

  • Finally UNet, the highest level PyTorch module in our model. It assembles the model using 3 EncoderBlocks, a Bottleneck (DoubleConv) in the middle, 3 DecoderBlocks, and a Conv2d output layer. In this project we won't use an activation function - e.g. a sigmoid - as the output layer, because we are going to use a loss function that works with raw logits: BCEWithLogitsLoss. For that, we need the unprocessed output information. Logits are the raw output values of the model before converting them into probabilities with a sigmoid function. Since we only predict whether each pixel is car or background, this is a binary segmentation problem, so one output channel with BCEWithLogitsLoss is enough.

Now take a break and spot an important difference between the forward methods of the EncoderBlock and the DecoderBlock.

The DecoderBlock has two inputs: x and skip. The second one denotes the Skip Connections marked with gray lines in the figure above.

Intuition Hook

You can think of a skip connection as a sort of IKEA assembly instruction.

The first DecoderBlock gets the parts (features) from the Bottleneck layer, but doesn't really know what to do with them.

In our project, the EncoderBlock from the same level gives the instructions using the skip connection saying, "I can provide this kind of image structure to fit the parts into", but it can't help further. Therefore, the first DecoderBlock assembles the parts based on the information it has, and hands it to the second DecoderBlock.

The second DecoderBlock has the assembly from the first DecoderBlock as one part, but has to build it in a bigger assembly now. To do this, it gets the assembly instructions from the EncoderBlock on the same level using the skip connection, and so on.

That's it! We just built a baseline model!

Training the Model

In the next step, open and copy the code of train.py.
In this file we will prepare our training and validation sets using the images in data/train/images and specify how to run a training and validation for each epoch.

The code in this file is pretty easy to follow. Start the coding part with the main function and refer to the comments for the explanation of the code. Don't forget about the if __name__ == "__main__": [...] part, otherwise in the next step, the training won't work as expected.

About the IoU metric: practically we won't use the compute_iou(...) (Intersection over Union) metric for anything. It's only there to see the training progress, and to have a better understanding about what's happening.

We won't use it as a loss function, because it returns discrete values and therefore it isn't differentiable. Backward propagation only works with differentiable functions. However, this is something we luckily don't have to do manually: frameworks like Tensorflow and PyTorch automatically do the calculations for us for those (see the #backward propagation part in the train_one_epoch(...) method).

Exploration

From now on, you will see some exciting results.
Simply follow the instructions of the two notebooks, 01_exploration.ipnyb and 02_training_baseline.ipnyb and try to write the code yourself for practice. They will guide you through the training and evaluation part.

For the 01_exploration.ipnyb, expected result should look something like this:

and for 02_training_baseline.ipnyb, something like this:

For better understanding of the numbers in the results, here is some rule of thumb:

  • High train_loss and High val_loss means high bias. This means that the model is underfitted, failing to recognise patterns.
  • Low train_loss and High val_loss means high variance/overfitting. The model is learned the training data "by heart". It can't recognize new data, it's too sensitive to the noise in the training data.
  • Low train_loss and Low val_loss normally means a good fit. The model learned to recognise the patterns, but immune to the noise.

If you'd like to see Part 2 of Image Segmentation in the next episode, please

Summary

In this part we learned how to read the a COCO-style dataset and handle it with PyTorch. We built a simple model with BCEWithLogitsLoss loss function (criterion). After a quick sanity check by overfitting 1 and then 10 images, we were able to train our first U-Net model.

Most importantly, we now hopefully have an intuition about the basic concepts of a neural networks and can apply them in real life.

Foreshadowing

In the next episode, Part 2 of Image Segmentation project we will focus on making our model even better.

We will use a great tool to curate the dataset, learn how to tune hyperparameters and build a better model architecture that can also be trained more consistently.

Follow if you'd like to join!

Afterword

Thank you for following this series and for coding along with this project. I appreciate it.

If you are someone learning Computer Vision, please leave a feedback in the comments and share your experience. It will help me imrpoving the quality of my work and it will help others to get better.

If you are an expert in the field, your constructive feedback means the world to me. I started this series to showcase concepts through interesting projects and help others to understand concepts the way I did. This means that my projects can have mistakes or missing information.

Top comments (0)