How to plot a segmentation mask – Best Tutorial Python

In this article, I’ll share with you the functions I’ve designed to quickly draw an image segmentation mask in Python.

One of the major tasks of Computer Vision is image segmentation.

Segmentation consists in dividing an image into groups of pixels called “segments”.

To perform segmentation, you need to determine which pixels belong to a segment.

Usually, a segment represents an object in an image.

By segmenting all the pixels of an object, we can easily determine its location and contours.

The precision required for segmentation makes it both a complex and powerful task in object detection.

Here, I’ll introduce you to two functions whose aim is to :

  • draw a single segmentation mask
  • plot several segmentation masks simultaneously

The last function is an extension of the first, so even if you have several masks to plot, I suggest you follow the development of the first function carefully.

Note: In what follows, I’m assuming that you already own a segmentation masks for an image. If this is not the case, you can consult my article on SAM a state of the art Deep Learning model from Meta at this link to learn how to use it and obtain your own segmentation masks from an image.

Function for drawing a single segmentation mask on an image

To draw a segmentation mask, we’ll use open-cv and numpy libraries.

We assume we have two numpy arrays:

  • the image with a shape (w, h, 3)
  • the mask with shape (w, h, 3)

Here, the mask is a Boolean array. It indicates True if the pixel is in the segment, False otherwise.

The following is the function to obtain the segmented image:

import cv2
import numpy as np

def draw_mask(image, mask_generated) :
  masked_image = image.copy()

  masked_image = np.where(mask_generated.astype(int),
                          np.array([0,255,0], dtype='uint8'),
                          masked_image)

  masked_image = masked_image.astype(np.uint8)

  return cv2.addWeighted(image, 0.3, masked_image, 0.7, 0)

The function first creates a copy of the image. Then it applies the mask using the np.where().

First segmentation step in Python

The result is an image with a mask over the segmented object.

Unfortunately, the mask is opaque: we can’t see through it. The segmented object is therefore indistinguishable.

To improve visualization, we overlay the starting image on the masked image (by modifying their opacity) with the code cv2.addWeighted(image, 0.3, masked_image, 0.7, 0). We then obtain an image that preserves both the mask and the segmented object.

Second segmentation step in Python

The image can then be displayed as follows :

By the way, if your goal is to master Deep Learning - I've prepared the Action plan to Master Neural networks. for you.

7 days of free advice from an Artificial Intelligence engineer to learn how to master neural networks from scratch:

  • Plan your training
  • Structure your projects
  • Develop your Artificial Intelligence algorithms

I have based this program on scientific facts, on approaches proven by researchers, but also on my own techniques, which I have devised as I have gained experience in the field of Deep Learning.

To access it, click here :

GET MY ACTION PLAN

GET MY ACTION PLAN

Now we can get back to what I was talking about earlier.

segmented_image = draw_mask(image, mask)

cv2.imshow(segmented_image)

With the draw_mask function, we can draw a segmentation mask on our image.

Let’s see how to adapt this function to use it with multiple masks.

Function for drawing multiple segmentation masks on an image

Segmentation masks are often output of Deep Learning models. In this context, two types of results are possible:

  • masks contained in a list of tables
  • masks are contained in a list of dictionaries

Here, we will adapt the function developed in Part 1 to these two particular cases.

From a list of tables

Here’s the function for drawing several masks on an image from a list of arrays:

def draw_masks_fromList(image, masks_generated, labels, colors) :
  masked_image = image.copy()
  for i in range(len(masks_generated)) :
    masked_image = np.where(np.repeat(masks_generated[i][:, :, np.newaxis], 3, axis=2),
                            np.asarray(colors[int(labels[i][-1])], dtype='uint8'),
                            masked_image)

    masked_image = masked_image.astype(np.uint8)

  return cv2.addWeighted(image, 0.3, masked_image, 0.7, 0)

We’ve added to the basic function :

  • a for loop to browse the list of masks
  • two lists :
    • labels corresponding to the detected classes
    • colors corresponding to the colors to be used for each label

Next, we can display the image:

COLORS = [(255,0,0),(255,255,0)]

LABELS = [[0], [0]] # 0 for bird, 1 for sheep

segmented_image = draw_mask(image, masks_list, LABELS, COLORS)

cv2.imshow(segmented_image)

A segmented image is obtained from a list of arrays.

From a list of dictionaries

Here’s the function for drawing several masks on an image from a list of dictionaries:

def draw_masks_fromDict(image, masks_generated) :
  masked_image = image.copy()
  for i in range(len(masks_generated)) :
    masked_image = np.where(np.repeat(masks_generated[i]['segmentation'].astype(int)[:, :, np.newaxis], 3, axis=2),
                            np.random.choice(range(256), size=3),
                            masked_image)

    masked_image = masked_image.astype(np.uint8)

  return cv2.addWeighted(image, 0.3, masked_image, 0.7, 0)

The function is designed for dictionaries containing masks with the key named 'segmentation', and for use without prior classes.

Next, you can display the image:

segmented_image = draw_mask(image, masks_dict)

cv2.imshow(segmented_image)

A segmented image is obtained from a list of dictionaries.

As you can see, drawing segmentation masks on an image is relatively quick. However, you may also want to draw segmentation masks on videos.

In that case, you can access my free Object Detection Off-road course. In it, you’ll learn how to segment videos easily with SAM.

To access your free training, click here :

unlock my training !

One last word, if you want to go further and learn about Deep Learning - I've prepared for you the Action plan to Master Neural networks. for you.

7 days of free advice from an Artificial Intelligence engineer to learn how to master neural networks from scratch:

  • Plan your training
  • Structure your projects
  • Develop your Artificial Intelligence algorithms

I have based this program on scientific facts, on approaches proven by researchers, but also on my own techniques, which I have devised as I have gained experience in the field of Deep Learning.

To access it, click here :

GET MY ACTION PLAN

GET MY ACTION PLAN

Tom Keldenich
Tom Keldenich

Artificial Intelligence engineer and data enthusiast!

Founder of the website Inside Machine Learning

Leave a Reply

Your email address will not be published. Required fields are marked *

This page will not stay online forever

Enter your email to receive for free

The PANE method for Deep Learning

* indicates required

 

You will receive one email per day for 7 days – then you will receive my newsletter.
Your information will never be given to third parties.

You can unsubscribe in 1 click from any of my emails.



Entre ton email pour recevoir gratuitement
la méthode PARÉ pour faire du Deep Learning


Tu recevras un email par jour pendant 7 jours - puis tu recevras ma newsletter.
Tes informations ne seront jamais cédées à des tiers.

Tu peux te désinscrire en 1 clic depuis n'importe lequel de mes emails.