How to manually create a label in Matplotlib.pyplot

Question

I am trying to compare 2 figures as follows:

I using the function subplot() under the library Matplotlib.pylot to plot these 2 images, and the following are my codes:

def Plot_Labels(data,data_check):
  data = np.reshape(data,(414,507))
  #data = np.reshape(data,(414,507))
  shape = data.shape
  nan_mask = np.isnan(data_check)
  data_mask = data[~nan_mask]
  valid_data = np.empty(shape)
  valid_data[~nan_mask] = data_mask
  valid_data[nan_mask] = np.nan

  fig, axes = plt.subplots()
  im = axes.imshow(valid_data)
  bar = plt.colorbar(im)
  bar.set_label('Flood Hazard')
  plt.show()

However, the labels in these 2 figures are different. For example, both 4 and 5 gives the same yellow color in the 2 figures respectively. Therefore, is there any method to manually create label 5 in the second figure, so that the 2 figures become comparable?