automatyczny_kelner/TreeConcept.py
2023-05-26 01:29:28 +02:00

38 lines
1.6 KiB
Python

from sklearn import tree
import pandas as pd #for manipulating the csv data
import numpy as np
import graphviz
import os
os.environ["PATH"] += os.pathsep + 'C:/Program Files (x86)/Graphviz/bin/'
#importing the dataset from the disk
train_data_m=np.genfromtxt("dataset/converted_dataset.csv", delimiter=",",skip_header=1);
# Separate the attributes and labels
X_train = [data[:-1] for data in train_data_m]
y_train = [data[-1] for data in train_data_m]
# Create the decision tree classifier using the ID3 algorithm
clf = tree.DecisionTreeClassifier(criterion='entropy')
#clf = tree.DecisionTreeClassifier(criterion='gini')
# Train the decision tree on the training data
clf.fit(X_train, y_train)
# Visualize the trained decision tree
tree_text = tree.export_text(clf,feature_names=['Battery Charge', 'Fullness', 'Ready orders', 'Waiting tables','Availability', 'Cleanliness', 'Error'])
with open('decision_tree.txt', 'w') as f:
f.write(tree_text) # Save the visualization as a text file
dot_data = tree.export_graphviz(clf, out_file=None, feature_names=['Battery Charge', 'Fullness', 'Ready orders', 'Waiting tables','Availability', 'Cleanliness', 'Error'], class_names=['NO', 'YES'], filled=True,rounded=True)
graph = graphviz.Source(dot_data)
graph.render("decision_tree") # Save the visualization as a PDF file
# Test the decision tree with a new example
#Battery Charge,Fullness,Ready orders,Waiting tables,Availability,Cleanliness,Error
new_example = [2, 0, 1, 1, 1 ,2, 0]
predicted_label = clf.predict([new_example])
if predicted_label[0]>0:
result="YES"
else:
result="NO"
print("Predicted Label:", result)