How to add legend to imshow() in matplotlib

Question

I am using matplotlib

In plot() or bar(), we can easily put legend, if we add labels to them. but what if it is a contourf() or imshow()

I know there is a colorbar() which can present the color range, but it is not satisfied. I want such a legend which have names(labels).

For what I can think of is that, add labels to each element in the matrix, then ,try legend(), to see if it works, but how to add label to the element, like a value??

in my case, the raw data is like:

1,2,3,3,4
2,3,4,4,5
1,1,1,2,2

for example, 1 represents 'grass', 2 represents 'sand', 3 represents 'hill'... and so on. imshow() works perfectly with my case, but without the legend.

my question is:

1. Is there a function that can automatically add legend, for example, in my case, I just have to do like this: someFunction('grass','sand',...)

2. If there isn't, how do I add labels to each value in the matrix. For example, label all the 1 in the matrix 'grass', labell all the 2 in the matrix 'sand'...and so on.

Thank you!

Edit:

Thanks to @dnalow, it's smart really. However, I still wonder if there is any formal solution.

Answer 1

I quote here a solution to a similar question, in case someone is still interested:

I suppose putting a legend for all values in a matrix only makes sense if there aren't too many of them. So let's assume you have 8 different values in your matrix. We can then create a proxy artist of the respective color for each of them and put them into a legend like this

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

# create some data
data = np.random.randint(0, 8, (5,5))
# get the unique values from data
# i.e. a sorted list of all values in data
values = np.unique(data.ravel())

plt.figure(figsize=(8,4))
im = plt.imshow(data, interpolation='none')

# get the colors of the values, according to the 
# colormap used by imshow
colors = [ im.cmap(im.norm(value)) for value in values]
# create a patch (proxy artist) for every color 
patches = [ mpatches.Patch(color=colors[i], label="Level {l}".format(l=values[i]) ) for i in range(len(values)) ]
# put those patched as legend-handles into the legend
plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0. )

plt.grid(True)
plt.show()

Answer 2

You could use matplotlib.pylab.text to add text to your plot and customize it to look like a legend

For example:

import numpy as np
import matplotlib.cm as cm
import matplotlib.pylab as plt

raw_data = np.random.random((100, 100))
fig, ax = plt.subplots(1)
ax.imshow(raw_data, interpolation='nearest', cmap=cm.gray)
ax.text(5, 5, 'your legend', bbox={'facecolor': 'white', 'pad': 10})
plt.show()

which gives you following

You can check out the matplotlib documentation on text for more details matplotlib text examples