Trashmaster/Untitled.ipynb

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]