SzybciorSmartTraktor/tree.py

30 lines
1.2 KiB
Python
Raw Normal View History

import matplotlib.image as pltimg
import matplotlib.pyplot as plt
import os
import pandas
import pickle
import pydotplus
from sklearn import tree
from sklearn.tree import DecisionTreeClassifier
def treelearn():
if os.path.exists("assets/tree.pkl"):
dtree = pickle.load(open(os.path.join('assets', "tree.pkl"), "rb"))
else:
df = pandas.read_csv(os.path.join('assets/data', 'data.csv'))
columns = ['Fuel','Water','Fertalizer','Carrots','Potatoes','Wheat','X','Y','seeds']
x = df[columns]
y = df['back to station']
dtree = DecisionTreeClassifier()
dtree = dtree.fit(x, y)
pickle.dump(dtree, open(os.path.join('assets', "tree.pkl"), "wb"))
data = tree.export_graphviz(dtree, out_file=None, feature_names=columns)
graph = pydotplus.graph_from_dot_data(data)
graph.write_png(os.path.join('assets', 'mytree.png'))
img = pltimg.imread(os.path.join('assets', 'mytree.png'))
imgplot = plt.imshow(img)
plt.show()
return dtree
def make_decision(tree, Fuel,Water,Fertalizer,Carrots,Potatoes,Wheat,X,Y,seeds):
decision = tree.predict([[Fuel, Water, Fertalizer, Carrots, Potatoes, Wheat, X, Y, seeds]])
return decision