DecisionTree save and load from file

This commit is contained in:
Łukasz 2021-05-31 21:05:42 +02:00
parent b3197dccd2
commit 01dae9185b
5 changed files with 10 additions and 4 deletions

View File

@ -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()

View File

@ -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)

View File

@ -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:

BIN
treedata/DecisionTree.pkl Normal file

Binary file not shown.

View File

Before

Width:  |  Height:  |  Size: 122 KiB

After

Width:  |  Height:  |  Size: 122 KiB