Spacing/Outline between pixels in px.imshow


I’m currently trying to create a heatmap using Plotly in order to be implemented into Dash and decided to convert my data into a pivot table to be represented as a heatmap using px.imshow.
Reason I decided to choose px.imshow instead of go.heatmap is because it has the option to make the pixel/cell sizes consistent no matter how the figure margins or number of data represented, which seems to be something go.heatmap isn’t capable of.
The issue with not using go.heatmap is that I seem to be abandoning the convenient xgap and ygap options to make it easier to differentiate between similar looking pixels.

This is my current code and output using px.imshow:

import as px

h = len(count_pivot.index.tolist()) * 30
w = len(count_pivot.columns.tolist()) * 30

fig = px.imshow(count_pivot, x=count_pivot.columns, y=count_pivot.index)
fig.update_xaxes(tickmode='linear', showgrid=False)
fig.update_yaxes(tickmode='linear', showgrid=False)
    'plot_bgcolor': 'rgba(255, 255, 255, 1)',
    'paper_bgcolor': 'rgba(255, 255, 255, 1)',
}, height=h, width=w)

What I want using xgap and ygap in go.heatmap:

Is there a way to incorporate this same functionality into px.imshow either as gaps/paddings between pixels or as simple white outlines around each pixel?

Thank you!

Never mind! I simply just needed to run a couple for loops to utilize fig.add_shape in order to create the rectangles necessary to create the grid.