Writing a tiff with colormap from rasterio

I'm writing a single band tif using the code below. What I'm most interested in doing is writing a colormap for band 0 along with the tiff.

The code below is iterating through some other files and is successfully writing out each band of each input. I would like to be able to apply a colormap to each output. Below I'm trying to ask for viridis to be limited to 5 categories.

I would like to avoid explicitly avoid importing GDAL - and this seems entirely possible.

So my question is:
How do I write out the file with the viridis colormap applied to it? If necessary, apply on band[0]

import xarray as xr
import rioxarray
import cftime
import os
import datetime
from rasterio.enums import Resampling
import matplotlib.pyplot as plt

ds_input = "inputfile.nc"
operatingdir = '\\'
filename = ds_input[:-3]

#determine the maximum number of bands in the raster using the band name in position 0 of the common_key list
bandcount = 7

#create a colormap for the output raster, with 5 colour bands
cmap = plt.cm.get_cmap('viridis', 5)

#use a for loop to output each day of the forecast
for i in range(0, bandcount):
    #Convert ds_netcdf_interp to raster using the max fbi values
    ds_input.var_name.rio.to_raster('{}\\Output\\{}_number{}.tif'.format(operatingdir, filename, i), 
        windowed=True, 
        compress='lzw', 
        dtype='int16', 
        nodata=-9999, 
        overwrite=True, 
        tiled=True, 
        cloud_optimized=True, 
        driver='COG', 
        colormap=cmap
    )

You can write a colormap to a tiff file using rasterios write_colormap(). Here's a short example on how to use it from the documentation:

import rasterio

with rasterio.Env():

    with rasterio.open('tests/data/shade.tif') as src:
        shade = src.read(1)
        meta = src.meta

    with rasterio.open('/tmp/colormap.tif', 'w', **meta) as dst:
        dst.write(shade, indexes=1)
        dst.write_colormap(
            1, {
                0: (255, 0, 0, 255),
                255: (0, 0, 255, 255) })
        cmap = dst.colormap(1)
        # True
        assert cmap[0] == (255, 0, 0, 255)
        # True
        assert cmap[255] == (0, 0, 255, 255)

Basically write_colormap takes the band number as the first argument and then a dictionary of raster value as keys and their corresponding RGBA values as values.