DecisionTree save and load from file

createDecisionTree.py

@@ -4,6 +4,7 @@ import pydotplus
 from sklearn.tree import DecisionTreeClassifier
 import matplotlib.pyplot as plt
 import matplotlib.image as pltimg
+import joblib
 
 df = pandas.read_csv("treedata\\data.csv")
 
@@ -38,12 +39,14 @@ y = df['Decision']
 
 dtree = DecisionTreeClassifier()
 dtree = dtree.fit(X, y)
+joblib.dump(dtree, 'treedata\\DecisionTree.pkl')
+
 
 data = tree.export_graphviz(dtree, out_file=None, feature_names=features)
 graph = pydotplus.graph_from_dot_data(data)
-graph.write_png('mydecisiontree.png')
+graph.write_png('treedata\\mydecisiontree.png')
 
-img = pltimg.imread('mydecisiontree.png')
+img = pltimg.imread('treedata\\mydecisiontree.png')
 imgplot = plt.imshow(img)
 plt.show()
 

createTreeData.py

@@ -119,4 +119,4 @@ filename = "treedata\\data.csv"
 with open(filename, 'w') as csvfile:
     writer = csv.DictWriter(csvfile, fieldnames=fields)
     writer.writeheader()
-    writer.writerows(dict)
+    writer.writerows(dict)

main.py

@@ -7,11 +7,11 @@ import WeatherConditions
 import convertToPrediction
 import drawUI
 import Graphsearch as graph
-from decisiontree import dtree
 from Tractor import Tractor
 from TractorLoad import TillageUnit
 from constants import *
 from manualSteering import manualSteeringDriver
+import joblib
 
 pygame.init()
 
@@ -58,6 +58,9 @@ while working:
         weather, day_time, temperature, wind, humidy = WeatherConditions.checkConditions()
 
     question = [convertToPrediction.convert(field.state, weather, day_time, temperature, wind, humidy)]
+
+    dtree = joblib.load('treedata\\DecisionTree.pkl')
+
     decision = dtree.predict(question)
 
     if tractor.autodrive: