skeleton_graph.py 文件源码

def plot_nx_graph_3d(graph, radii = None,  colormap='jet', line_width = 2, opacity=.9):
  """Plot a 3d graph of the skeleton

  Arguments:
    radii: radii of the edges used in color code, if None uniform color

  Returns:
    mayavi scence
  """
  # get graph positions
  g2 = nx.convert_node_labels_to_integers(graph, label_attribute = 'xyz');
  xyz = np.array([x['xyz'] for x in g2.node.values()], dtype = 'int32');

  # scalar colors
  if radii is not None:
    scalars = np.array([radii[tuple(x)] for x in xyz], dtype = 'float32');
  else:
    #scalars = np.arange(5, xyz.shape[0]+5);
    scalars = np.ones(xyz.shape[0], dtype = 'float32');

  #pts = mlab.points3d(xyz[:,0], xyz[:,1], xyz[:,2],
  #                    scalars,
  #                    scale_factor=node_size,
  #                    scale_mode='none',
  #                    colormap=graph_colormap,
  #                    resolution=20)

  pts = mlab.pipeline.scalar_scatter(xyz[:,0], xyz[:,1], xyz[:,2], scalars)

  pts.mlab_source.dataset.lines = np.array(g2.edges(), dtype = 'int32')
  pts.update()    

  #tube = mlab.pipeline.tube(pts, tube_radius=edge_size)
  #lab.pipeline.surface(tube, color=edge_color)

  lines = mlab.pipeline.stripper(pts);
  mlab.pipeline.surface(lines, colormap = colormap, line_width = line_width, opacity = opacity)

  if radii is not None:
      mlab.colorbar(orientation = 'vertical', title='Radius [pix]');    

  mlab.axes()

  return lines