Python script to add a colored border to a semantic segmentation map.

add-semseg-border.py

import argparse
import os

import imageio
from skimage import measure
import numpy as np

def parse_args():
    parser = argparse.ArgumentParser(description='Script to add a color border to semantic images.')
    parser.add_argument('image', type=str, default="", help="Input image to convert")
    parser.add_argument('--color', type=str, default='#ffffff', help='Color of the border as a 6 digit hex.')
    parser.add_argument('--suffix', type=str, default='_border', help='Suffix to add to the output file name.')
    parser.add_argument('--output_dir', type=str, default=None, help='Set a separate output directory.')
    args = parser.parse_args()

    return args

def main(args):
    image = imageio.imread(args.image)
    colors = np.unique(image.reshape(-1,image.shape[2]), axis=0)

    # create a greyscale image, where each color represents a different grey value
    gray_image = np.zeros(image.shape[:2]).astype(np.uint8)
    border_mask = np.zeros(image.shape[:2]).astype(np.bool)

    for idx, color in enumerate(colors):
        mask = np.all(image == color, axis=-1)
        gray_image[mask] = idx

    total_border_pixel = 0

    for idx in range(len(colors)):
        contours = measure.find_contours(gray_image, idx+0.5, positive_orientation='low')
        contours.extend( measure.find_contours(gray_image, idx+0.5, positive_orientation='high') )
        for contour in contours:
            border_mask[contour[:,0].astype(np.int),contour[:,1].astype(np.int)] = True
            total_border_pixel += contour.shape[0]

    border_pixel = np.ones(image.shape[2]) * 255
    border_pixel[0] = int(args.color[1:3],16)
    border_pixel[1] = int(args.color[3:5],16)
    border_pixel[2] = int(args.color[5:7],16)
    image[border_mask] = border_pixel

    filename, extension = os.path.splitext(args.image)
    output_filename = filename + args.suffix + extension

    # change the output directory is specified
    if args.output_dir is not None:
        output_filename = os.path.join(args.output_dir, os.path.basename(output_filename))
    
    imageio.imsave(output_filename, image)


if __name__ == '__main__':
    args = parse_args()
    main(args)