Pass custom colors to imshow

Hi

Thanks for reading and the help.

Iโ€™m trying to generate an image with custom colors. Image is 2D and ideally Iโ€™d like to pass a dictionary mapping values to colors, something like this

import pandas as pd
import plotly.express as px
from datetime import time

im_2d = np.random.randint(0,4, (24, 7))

colors = {
    0: (0,0,0),
    1: (255, 0, 0),
    2: (0, 255, 0),
    3: (0, 0, 255)
}
fig1 = px.imshow(
          im_2d,
          color_diescrete_map = colors, # Or something along this line
          x=pd.date_range(start='2020-01-01', end='2020-01-07'),
          y = [time(t, 0) for t in range(24)]
         )

Since that doesnโ€™t seem to be implemented (please correct me if Iโ€™m wrong), Iโ€™m trying to generate a 3D array where I put the RGB colors from my dictionary, but when doing that, imshow throws an exception (related to my fancy indexing), please execute this to reproduce. The exception is related to the fancy indexing since removing them causes no exception

im_3d = np.random.randint(0, 4, (24, 7, 3))
fig2 = px.imshow(im_3d,
          x=pd.date_range(start='2020-01-01', end='2020-01-07'),
          y = [time(t, 0) for t in range(24)]
         )

Thanks again for the help