30 lines
1.2 KiB
Python
30 lines
1.2 KiB
Python
|
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
|