2.4 KiB
2.4 KiB
from sklearn.tree import DecisionTreeClassifier, export_text, plot_tree
dataset = pandas.read_csv('/Users/mac/Desktop/tree_dataset.csv', sep=";")
decisions = ["decision"]
attributes = ["season", "trash_type", "mass", "space", "trash_mass"]
x = dataset[attributes]
y = dataset[decisions]
decision_tree = DecisionTreeClassifier()
decision_tree.fit(x.values, y.values)
decision = decision_tree.predict(
[[5, 3 , 3, 1, 2]])
print(decision)
[1]