65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
import joblib
|
|
import matplotlib.pyplot as plt
|
|
import pandas
|
|
from sklearn.tree import DecisionTreeClassifier, export_text, plot_tree
|
|
from Sectors_types import Sectors_types
|
|
|
|
|
|
decisions = ["decision"]
|
|
attributes = ["marking", "weekend", "c_popularity", "payment_delay", "payed", "shipping_method"]
|
|
|
|
|
|
class DecisionTree:
|
|
|
|
def __init__(self, marking, weekend, c_popularity, payment_delay, payed, shipping_method):
|
|
self.decision = self.decision(marking, weekend, c_popularity, payment_delay, payed, shipping_method)
|
|
|
|
def tree(self):
|
|
dataset = pandas.read_csv('./decision_tree/csv_file.csv')
|
|
x = dataset[attributes]
|
|
y = dataset[decisions]
|
|
decision_tree = DecisionTreeClassifier()
|
|
decision_tree = decision_tree.fit(x.values, y)
|
|
|
|
return decision_tree
|
|
|
|
# return decision made from tree and attributes
|
|
def decision(self, marking, weekend, c_popularity, payment_delay, payed, shipping_method):
|
|
decision_tree = self.tree()
|
|
decision = decision_tree.predict(
|
|
[[marking, weekend, c_popularity, payment_delay, payed, shipping_method]])
|
|
if decision == 1:
|
|
decision = Sectors_types.fragile
|
|
elif decision == 2:
|
|
decision = Sectors_types.normal
|
|
elif decision == 3:
|
|
decision = Sectors_types.shipping_tomorrow
|
|
elif decision == 4:
|
|
decision = Sectors_types.shipping_today
|
|
return decision
|
|
|
|
def tree_as_txt(self, decision_tree):
|
|
with open('tree.txt', "w") as file:
|
|
file.write(export_text(decision_tree))
|
|
|
|
def tree_to_png(self, decision_tree):
|
|
fig = plt.figure()
|
|
plot_tree(decision_tree, feature_names=attributes, filled=True)
|
|
plt.title("Decision tree")
|
|
# plt.show()
|
|
fig.savefig('tree.png')
|
|
|
|
# def tree_to_structure(self,decision_tree):
|
|
# joblib.dump(decision_tree, 'tree_model')
|
|
#
|
|
# def tree_from_structure(self, file):
|
|
# return joblib.load(file)
|
|
#
|
|
#
|
|
#
|
|
# drzewo = tree()
|
|
# # tree_as_txt(drzewo)
|
|
# tree_to_png(drzewo)
|
|
# # tree_to_structure(drzewo)
|
|
|