draw.py 文件源码

def draw(m, name, extra=None):
  FIG.clf()

  matrix = m
  orig_shape = np.shape(matrix)
  # lose the channel shape in the end of orig_shape
  new_shape = orig_shape[:-1] 
  matrix = np.reshape(matrix, new_shape)
  ax = FIG.add_subplot(1,1,1)
  ax.set_aspect('equal')
  plt.imshow(matrix, interpolation='nearest', cmap=plt.cm.gray)
  # plt.imshow(matrix, interpolation='nearest', cmap=plt.cm.ocean)
  plt.colorbar()

  if extra != None:
    greens, reds = extra
    grn_x, grn_y, = greens
    red_x, red_y = reds
    plt.scatter(x=grn_x, y=grn_y, c='g', s=40)
    plt.scatter(x=red_x, y=red_y, c='r', s=40)
#  # put a blue dot at (10, 20)
#  plt.scatter([10], [20])
#  # put a red dot, size 40, at 2 locations:
#  plt.scatter(x=[3, 4], y=[5, 6], c='r', s=40)
#  # plt.plot()

  plt.savefig(name)