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().

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.

The image can then be displayed as follows :

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 !