Beispiel #1
0
def draw_tree(tree, path):
    dot_data = StringIO()

    dot_data.seek(0)  # to reset graph?
    dot_data.truncate(0)  # to reset graph?

    export_graphviz(tree,
                    out_file=dot_data,
                    feature_names=tree.feature_names,
                    filled=True,
                    impurity=False,
                    proportion=True,
                    rounded=True,
                    special_characters=True,
                    node_ids=True)
    graph = pydotplus.graph_from_dot_data(dot_data.getvalue())
    open(path, 'wb').write(Image(graph.create_png()).data)
Beispiel #2
0
test_prediction_rate = get_correct_ratio(test_results, test_survived)

print('{} : {}'.format("test_prediction_rate", test_prediction_rate))

dot_data = StringIO()

export_graphviz(titanic_tree,
                out_file=dot_data,
                filled=True,
                rounded=True,
                special_characters=True,
                feature_names=titanic_data_features.columns.values,
                class_names=["Died", "Survived"])

graph = pydotplus.graph_from_dot_data(dot_data.getvalue())

graph.write_png('tree.png')

png_str = graph.create_png(prog='dot')

# treat the dot output string as an image file
sio = StringIO()
sio.write(png_str)
sio.seek(0)
img = mpimg.imread(sio)

# plot the image
imgplot = plt.imshow(img, aspect='equal')
plt.show(block=False)