How to color a label image in python

label_to_color.py

#!/usr/bin/python
######################################################################
#
#  An example of how to color code a gray level label image in
#  python.
#
#  2021-05-05 Wed
#  Dov Grobgeld <dov.grobgeld@gmail.com>
######################################################################

import numpy as np
from PIL import Image

# Create a color table
MyTable = np.array([[0,0,0],   # 0
                    [255,0,0], # 1
                    [0,255,0], # 2
                    [0,0,255], # 3
                    [255,255,0], # 4
                    [255,0,255], # 5
                    [0,255,255], # 6
                    [0,128,255], # 7
                    ] + [[0,0,0]]*(256-8),
                    dtype=np.uint8)

                             
# Create a 2d array
a = np.array([[1,1,1,1,1],
              [1,2,2,2,1],
              [1,2,3,2,1],
              [1,2,2,2,1],
              [1,1,1,1,1]],dtype=np.uint8)

# Color it
a_color = np.take(MyTable,
                  a,axis=0)
# Save as image
Image.fromarray(a_color).save('label-image.png')