Trashmaster/decision_tree/decisionTree.py
aliaksei-kudravets f258750c1c fix path
2022-05-13 12:47:08 +02:00

59 lines
1.4 KiB
Python

import joblib
import matplotlib.pyplot as plt
import pandas
from sklearn.tree import DecisionTreeClassifier, export_text, plot_tree
decisions = ["decision"]
attributes = ["season", "trash_type", "mass", "space", "trash_mass"]
# return tree made from attributes
def tree():
dataset = pandas.read_csv('decision_tree/tree_dataset.csv', sep=";")
x = dataset[attributes]
y = dataset[decisions]
decision_tree = DecisionTreeClassifier()
decision_tree = decision_tree.fit(x.values, y.values)
return decision_tree
# return decision made from tree and attributes
def decision(decision_tree, season, trash_type, mass, space, trash_mass):
decision = decision_tree.predict(
[[season, trash_type , mass, space, trash_mass]])
return decision
'''
we shall save output of our decision tree. It is possible for a few ways:
txt, png or structure
'''
def tree_as_txt(decision_tree):
with open('./decision_tree/tree_as_txt.txt', "w") as file:
file.write(export_text(decision_tree))
def tree_to_png(decision_tree):
plt.figure()
plot_tree(decision_tree, feature_names=attributes, filled=True)
plt.title("Decision tree")
plt.show()
def tree_to_structure(decision_tree):
joblib.dump(decision_tree, './decision_tree/tree_model')
def tree_from_structure(file):
return joblib.load(file)
#drzewo = tree()
#tree_as_txt(drzewo)
#tree_to_png(drzewo)
#tree_to_structure(drzewo)