How to color cells in a plt plot based on value?

Question

Im trying to use pylot to illustrate a layout in a robotic path planning RL concept. I've created the layout as below, but I cannot get the cells to color based on the value of the cell. I've used an np.array to generate the image and pyplot, but matlib colors it based on a std heatmap.

I have tried the following just to color all cells as a start:

ax.set_facecolor('red')

But it seem to color the plot behind the ax.imshow, rendering it hidden on the illustration.

The entire .py:

import numpy as np
import matplotlib
import matplotlib.pyplot as plt

layout = np.array([
 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
 [1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
 [1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1],
 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1],
 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1],
 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]
)


fig, ax = plt.subplots()
im = ax.imshow(layout)


ax.set_xticks(np.arange(len(layout[1])))
ax.set_yticks(np.arange(len(layout)))

# Rotate the tick labels and set their alignment.
plt.setp(ax.get_xticklabels(), rotation=45, ha="right",
         rotation_mode="anchor")

# Loop over data dimensions and create text annotations.
for i in range(len(layout)):
    for j in range(len(layout[1])):
        text = ax.text(j, i, layout[i, j],
                       ha="center", va="center", color="w")
        ax.set_facecolor('red')

ax.set_title("Showroom layout")
fig.tight_layout()

plt.show()

I know the set_facecolor right now will color the entire plot red, but im still debugging. In the future it will color cells with value "1" gray and "0" white. Or alike.

I appreciate your help and input! :)