plot_utils.py 文件源码

def lblshow(label_img, labels_str=None, f=None, ax=None, cmap=None, *args, **kwargs):
  ''' display a labeled image with associated legend

  Parameters
  ----------
  label_img : labeled image [nrows, ncols] = numpy.array.shape
  labels_str : a complete list of labels
  f : (optional) a figure handle
  cmap : the color of each label (optional). like a list of colors, e.g.,
      ['Red','Green',...] or a matplotlib.colors.ListedColormap)
  '''

  if labels_str is None:
    labels_str = [str(i) for i in np.unique(label_img)]

  if ax is None:
    if f is None:
      f,ax = plt.subplots(1,1)
      f.set_size_inches(9,6)
    else:
      ax = f.gca()
  elif f is None:
    f = ax.get_figure() 


  nlabels = len(labels_str)
  if type(cmap) is mpl.colors.ListedColormap:
    pass
  elif hasattr(cmap, '__iter__'):
    if not kwargs.has_key('norm'):
      bounds = range(0,len(cmap)+1)
      kwargs['norm'] = mpl.colors.BoundaryNorm(bounds, len(cmap)) # HACKY
    cmap = mpl.colors.ListedColormap(cmap)
  elif cmap is None:
    colors = mpl.cm.spectral(np.linspace(0, 1, nlabels))
    cmap = mpl.colors.ListedColormap(colors)
  else:
    assert False, 'invalid color map'


  im = ax.imshow(label_img, cmap=cmap, *args, **kwargs); ax.axis('off')

  # create an axes on the right side of ax. The width of cax will be 5%
  # of ax and the padding between cax and ax will be fixed at 0.05 inch.
  divider = make_axes_locatable(ax)
  cax = divider.append_axes("right", size="5%", pad=0.05)
  cbar = plt.colorbar(im, cax=cax)

  cbar.ax.get_yaxis().set_ticks([])
  for j, lab in enumerate(labels_str):
    cbar.ax.text(1.3, float(2 * j + 1) / (nlabels*2), lab, ha='left', va='center')

  return f