How to assign particular color to each pixel in image

I have a single band raster tif image with 5 values i.e 4,3, 2, 1, 0 .

This code displays image with random color assigned to each pixel value

import rasterio
from rasterio.plot import show
import matplotlib.pyplot as plt

img_data = rasterio.open(img)
fig, ax = plt.subplots(figsize=(10,10))    
show(img_data)

How can I display this image by assigning particular color to each image (E.g. red: 4,blue:3, green:2, black:1, white: 0). I came across colormap option but I am not able to display it as desired.

Solution

I recently solved this, so here is what worked for me. Import your raster as a numpy array

import rasterio as rio
with rio.open(path_img) as img:
    img_data = img.read()

First you need to create a colormap:

from matplotlib import colors
#since your values are from 0 to 4, you want the color bins to be at 0.5 increments
levels = [-0.5,0.5,1.5,2.5,3.5,4.5]
clrs = ['white','black','green','blue','red'] 
cmap, norm = colors.from_levels_and_colors(levels, clrs)

Next you want to create your plot:

fig, ax = plt.subplots(figsize=(15,15))
raster_plot = show(img_data,ax=ax,cmap=cmap,norm=norm)

Now you create the colorbar:

#get the image values
im = raster_plot.get_images()[0]
#create a colorbar
cbar = fig.colorbar(im, ax=ax)
cbar.ax.get_yaxis().set_ticks([])
#add text to the middle of the colorbar values
for j, lab in enumerate(['0','1','2','3','4']):
    cbar.ax.text(2, j, lab, ha='center', va='center')