Categorical legend for heatmap to plot bin information of wafer

Hi, I want to create a interactive wafer map for the bin information (e.g. bin 1 good chips, bin 2 bad chips, bin 3 not tested). It seems for me that px.imshow() (python - plotly marker size relative to data to plot interactive wafer map - Stack Overflow) is the way to go but I cant figure out how to get a categorical legend and color mapping. I want to use the feature of double clicking in the legend to turn on and of traces. It seems to be possible (github.com/plotly/plotly.js/issues/2642 , Categorical choropleths with legends ยท Issue #3468 ยท plotly/plotly.js ยท GitHub) but I dont know how.