Decision trees visualization

visualize_decision_tree.py

from sklearn.datasets import load_iris
from sklearn import tree
import numpy as np


iris = load_iris()

features = iris.data
labels = iris.target

test_indexes = [2, 52, 102]

# training data
train_features = np.delete(features, test_indexes, axis=0)
train_labels = np.delete(labels, test_indexes)
##print(len(train_features), len(train_labels))


# testing data
test_features = features[test_indexes]
test_labels = labels[test_indexes]
##print(len(test_features), len(test_labels))

clf = tree.DecisionTreeClassifier()
clf.fit(train_features, train_labels)

ans = iris.target_names[clf.predict(test_features)]

print('The test result for {} is {} respectively'.format(test_indexes, ans))

#viz code
from sklearn.externals.six import StringIO
import pydot

dot_data = StringIO()
tree.export_graphviz(clf,
                     out_file=dot_data,
                     feature_names=iris.feature_names,
                     class_names=iris.target_names,
                     filled=True, rounded=True,
                     impurity=False)

graph = pydot.graph_from_dot_data(dot_data.getvalue())
graph[0].write_pdf('iris.pdf')