Python: putting segments on top of the input image

Assume you are working on the image /video segmentation problem. Sooner or later a day is coming when you need to present your (hopefully good) results to other people on a conference, project meeting, in a paper, whatever. Usually segmentation results are represented by a matrix of integers standing for region labels which are supposed to be unique. A colored representation of segments is derived by assignment of unique colors based on region labels. In this post I will show you a simple trick how you can impress the audience with your segmentation results without any changes in your segmentation method. Using Alpha Compositing you can combine an input image with colored segments as it’s shown below:

For this I will use Python, NumPy, and Python Image Library (PIL).

import click
from PIL import Image


@click.command(help='Create image with segmentation overlay')
@click.argument('image_path', type=click.Path(exists=True))
@click.argument('segments_png_path', type=click.Path(exists=True))
def main(image_path: str, segments_png_path: str) -> None:
    image = Image.open(image_path)
    segments = Image.open(segments_png_path)

    map_mesh = Image.new('RGBA', image.size, (0, 0, 0, 100))
    image.paste(segments, (0, 0), map_mesh)
    image.save('output.png')


if __name__ == '__main__':
    # pylint: disable=no-value-for-parameter
    main()

First of all we need a function generating colors for segment labels given by your segmentation technique. It can be done as follows:

import random

import click
import numpy as np
from PIL import Image

SEGMENTS_N = 3000

Color = tuple[int, int, int]


def generate_colors() -> list[Color]:
    def _random_number() -> int:
        return random.randint(0, 255)

    def _random_color() -> Color:
        return _random_number(), _random_number(), _random_number()

    return [_random_color() for _ in range(SEGMENTS_N)]

Here we don’t check the uniqueness of generated colors, but such a check can be done easily. The following function returns colored segments in form of a 2-dimensional numpy array:

def colored_segments(segments: np.ndarray, colors: list[Color]) -> np.ndarray:
    height, width = segments.shape
    segments_image = np.zeros([height, width, 3], dtype=np.uint8)
    segment_ids = np.unique(segments)
    for id_ in segment_ids:
        segments_image[:, :, :][(segments == id_)] = colors[id_]

    return segments_image

Now we call both functions to get a colored representation of segments and put segments on top of the input image. For this we use paste function from the Python Image Library (PIL) which takes a transparency mask as an optional argument. The transparency value (here $140$ ) is up to you depending on which effect you want to achieve.

@click.command(help='Create image with segmentation overlay')
@click.argument('image_path', type=click.Path(exists=True))
@click.argument('segments_npy_path', type=click.Path(exists=True))
def main(image_path: str, segments_npy_path: str) -> None:
    image = Image.open(image_path)
    segments = np.load(segments_npy_path)

    colors = generate_colors()
    segments_mask = colored_segments(segments, colors)
    segments_mask = Image.fromarray(np.uint8(segments_mask))

    map_mesh = Image.new('RGBA', image.size, (0, 0, 0, 100))
    image.paste(segments_mask, (0, 0), map_mesh)
    image.save('output.png')


if __name__ == '__main__':
    # pylint: disable=no-value-for-parameter
    main()

Finally I would like to show you a video demonstrating segmentation of a video sequence. Final segments are shown here for each frame in conjunction with the original frame using the alpha compositing. The sequence was segmented by this segmentation technique. Enjoy!