用pydot绘制决策树

Zhu*_*arb 5 python graphviz decision-tree pydot

我已经训练了一个决定tree(Python字典),如下所示.现在我想用pydot绘制它.在定义树的每个节点(pydot图)时,我为它指定了一个唯一(和详细)的名称和一个简短的标签.

我的问题是,通过写入.png得到的结果,我看到了详细node names而不是node labels.

在这里听了@Martijn Pieters的回答.我不知道我错过了什么,有什么想法吗?

import pydot

tree= {'salary': {'41k-45k': 'junior', '46k-50k': {'department': {'marketing': 'senior', 'sales': 'senior', 'systems': 'junior'}}, '36k-40k': 'senior', '26k-30k': 'junior', '31k-35k': 'junior', '66k-70k': 'senior'}}

def walk_dictionaryv2(graph, dictionary, parent_node=None):
    '''
    Recursive plotting function for the decision tree stored as a dictionary
    '''

    for k in dictionary.keys():

        if parent_node is not None:

            from_name = parent_node.get_name().replace("\"", "") + '_' + str(k)
            from_label = str(k)

            node_from = pydot.Node(from_name, label=from_label)

            graph.add_edge( pydot.Edge(parent_node, node_from) )

            if isinstance(dictionary[k], dict): # if interim node


                walk_dictionaryv2(graph, dictionary[k], node_from)

            else: # if leaf node
                to_name = str(k) + '_' + str(dictionary[k]) # unique name
                to_label = str(dictionary[k])

                node_to = pydot.Node(to_name, label=to_label, shape='box')
                graph.add_edge(pydot.Edge(node_from, node_to))

                #node_from.set_name(to_name)

        else:

            from_name =  str(k)
            from_label = str(k)

            node_from = pydot.Node(from_name, label=from_label)
            walk_dictionaryv2(graph, dictionary[k], node_from)


def plot_tree(tree, name):

    # first you create a new graph, you do that with pydot.Dot()
    graph = pydot.Dot(graph_type='graph')

    walk_dictionaryv2(graph, tree)

    graph.write_png(name+'.png')


plot_tree(tree,'name')
Run Code Online (Sandbox Code Playgroud)

这是上面代码得到的(不需要的)输出:

在此输入图像描述

Mar*_*ers 4

您需要将您创建的节点显式添加到图中:

node_from = pydot.Node(from_name, label=from_label)
graph.add_node(node_from)
Run Code Online (Sandbox Code Playgroud)

node_to = pydot.Node(to_name, label=to_label, shape='box')
graph.add_node(node_to)
Run Code Online (Sandbox Code Playgroud)

否则渲染器将看不到名称。graph.add_node()在生成的文件中包含节点元数据.dot

添加这些graph.add_node()行后,结果是:

名称.png