Plotting multi-class semantic segmentation transparent overlays over RGB image

Question

I have the results of semantic segmentation masks (values between 0-1, requiring otsu thresholding to determine what's positive) which I'd like to plot directly on the RGB image with different random color per prediction class on an RGB image.

I used the following to plot a single mask with a single color. Is there a package or simple strategy to do that for multi-class?

fig, ax = plt.subplots(nrows=nrows, ncols=ncols, figsize=(5, 5))
  ax.imshow(image, cmap='gray')
  ax.axis('off')
  mask = (fused_mosaic[..., channel]*255).astype('uint8')
  ret3,th3 = cv2.threshold(mask,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
  fig, ax = image_show(full_im)
  ax.imshow(mask>ret3, alpha=0.3)

I’m looking for something like this, just simples without the boxes and labels. I tried using detectron2 (the package that generates this annotation in the example, but they require some odd metadata object which I can’t figure out).

Thanks

Answer 1

Scikit-image has a built-in label2rgb() function that colorises according to a label channel:

#!/usr/bin/env python3

from skimage import io
from skimage import color
from skimage import segmentation
import matplotlib.pyplot as plt

# URL for tiger image from Berkeley Segmentation Data Set BSDS
url=('http://www.eecs.berkeley.edu/Research/Projects/CS/vision/bsds/BSDS300/html/images/plain/normal/color/108073.jpg')

# Load tiger image from URL
tiger = io.imread(url)

# Segment image with SLIC - Simple Linear Iterative Clustering
seg = segmentation.slic(tiger, n_segments=30, compactness=40.0, enforce_connectivity=True, sigma=3)

# Generate automatic colouring from classification labels
io.imshow(color.label2rgb(seg,tiger))
plt.show()

Answer 2

With our conversation above, you have a 2D NumPy array of integer IDs where each element in this array determines the class ID of said pixel thus giving you a semantic segmentation output.

I would recommend you do this in three stages.

1. Create a RGB colour map that is of size N x 4 where N is the total number output classes in your segmentation. Therefore, each class i is assigned a RGBA colour pixel value that you would use to colour the output.

2. Flatten the input integer NumPy array so that it's a 1D NumPy array that we can use to index into (1)

3. Finally index into the RGB colour map in (1). This will create a (R x C) x 4 2D NumPy array which contains the output colour image mapping semantic labels to colours. Of course we need this back in the original input dimensions so reshape this array so that it becomes R x C x 4 for you to display. Finally, because we now have an alpha channel for the image, you can just display this on top of your original image.

---

Step #1 - Generate the colour map

matplotlib has a nice set of tools to generate this colour map for you. You can use the cm module from this. First, decide what colour map you'd like to use for your purposes. The full list of them can be found here: https://matplotlib.org/3.1.1/tutorials/colors/colormaps.html. I'll go with Viridis as that is the default currently being used in matplotlib.

Assuming that the total number of classes you have in your system is N, first generate the colour map, then create a linearly spaced array from 0 to 1 with N elements to uniformly create colours from the beginning to the end of this colour map. Also take note that this will generate a N x 4 colour map with the last column being the alpha channel. This is extremely important for later. Specifically, this method will colour any pixel with label 0 to belong to the lower end of the colour map and because this is a semantic segmentation output, label 0 should correspond to the background so we should set the alpha channel for this label to be 0 to be transparent. The rest of the colours we can set to your desired alpha, which is 0.3 in your code.

from matplotlib import cm
import numpy as np
N = ... # You define this here
colours = cm.get_cmap('viridis', N)  # Change the string from 'viridis' to whatever you want from the above link
cmap = colours(np.linspace(0, 1, N))  # Obtain RGB colour map
cmap[0,-1] = 0  # Set alpha for label 0 to be 0
cmap[1:,-1] = 0.3  # Set the other alphas for the labels to be 0.3

Step #2 - Take your semantic segmentation output and find the appropriate colours

This is straight forward. Assuming fused_mosaic is the 2D integer array we discussed earlier, flatten this array and index your colour map:

output = cmap[fused_mosaic.flatten()]

Step #3 - Reshape to the desired output

This again is straight forward:

R, C = fused_mosaic.shape[:2]
output = output.reshape((R, C, -1))

output will now contain your RGBA rendered image for each object that's in your semantic segmentation map. You can then finally use this and display this on top of your image. With your code, this would be:

fig, ax = image_show(full_im)  # Don't know what this does but it's from your code
ax.imshow(output)

---

To tie everything together, this is finally what I'd do:

## Step #1
from matplotlib import cm
import numpy as np
N = ... # You define this here
colours = cm.get_cmap('viridis', N)  # Change the string from 'viridis' to whatever you want from the above link
cmap = colours(np.linspace(0, 1, N))  # Obtain RGB colour map
cmap[0,-1] = 0  # Set alpha for label 0 to be 0
cmap[1:,-1] = 0.3  # Set the other alphas for the labels to be 0.3

## Step #2
output = cmap[fused_mosaic.flatten()]

## Step #3
R, C = fused_mosaic.shape[:2]
output = output.reshape((R, C, -1))

## Overlay
fig, ax = image_show(full_im)  # Don't know what this does but it's from your code
ax.imshow(output)