1 files changed, 17 insertions, 0 deletions
diff --git a/Assignment 3/packages/treeViewer.py b/Assignment 3/packages/treeViewer.py
new file mode 100755
index 0000000..eadcc28
--- /dev/null
+++ b/Assignment 3/packages/treeViewer.py
@@ -0,0 +1,17 @@
+from sklearn.externals.six import StringIO
+from sklearn import tree
+import pydot
+
+FNAMES = ['Fixed acidity', 'Volatile acidity', 'Citric acid', 'Residual sugar', 'Chlorides', 'Free sulfur dioxide', 'Total sulfur dioxide', 'Density', 'pH', 'Sulphates', 'Alcohol']
+
+def view(classifier):
+    """ Renders a graph representation of classifier, and
+        saves it to "MyTree.pdf" in the same folder
+        as the executing script.
+    """
+    tree_dot = StringIO()
+    tree.export_graphviz(classifier, out_file=tree_dot, feature_names=FNAMES)
+    graph = pydot.graph_from_dot_data(tree_dot.getvalue())
+    print "A"
+    graph.write_pdf("MyTree.pdf")
+    print "B"