diff --git a/machine_learning/decisionTree.py b/machine_learning/decisionTree.py index 17cc4e1..68a3950 100644 --- a/machine_learning/decisionTree.py +++ b/machine_learning/decisionTree.py @@ -1,13 +1,15 @@ import os from trainingData import TrainingData - +from sklearn import tree +import joblib +from sklearn.preprocessing import LabelEncoder, OneHotEncoder +import numpy as np def _read_training_data() -> TrainingData: - attributes: list[list[str]] = [] - classes: list[str] = [] - __location__ = os.path.realpath( - os.path.join(os.getcwd(), os.path.dirname(__file__))) - file = open(os.path.join(__location__, 'training_data.csv')) + attributes = [] + classes = [] + location = os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__))) + file = open(os.path.join(location, 'training_data.csv')) lines = file.readlines()[1:] file.close() for line in lines: @@ -19,5 +21,29 @@ def _read_training_data() -> TrainingData: classes.append(line_class) return TrainingData(attributes, classes) + trainning_data = _read_training_data() - \ No newline at end of file + +X = trainning_data.attributes +Y = trainning_data.classes + +le_shape = LabelEncoder() +le_flexibility = LabelEncoder() +le_color = LabelEncoder() + +le_shape.fit([x[0] for x in X]) +le_flexibility.fit([x[3] for x in X]) +le_color.fit([x[4] for x in X]) + +X_encoded = np.array([ + [le_shape.transform([x[0]])[0], x[1], x[2], le_flexibility.transform([x[3]])[0], le_color.transform([x[4]])[0]] + for x in X +]) + +encoder = OneHotEncoder(categories='auto', sparse=False) +X_encoded = encoder.fit_transform(X_encoded) + +model = tree.DecisionTreeClassifier() +model.fit(X_encoded, Y) + +joblib.dump(model, 'model.pkl') \ No newline at end of file diff --git a/machine_learning/model.pkl b/machine_learning/model.pkl new file mode 100644 index 0000000..f79ca17 Binary files /dev/null and b/machine_learning/model.pkl differ diff --git a/machine_learning/training_data.csv b/machine_learning/training_data.csv index b0d47b1..3e7951e 100644 --- a/machine_learning/training_data.csv +++ b/machine_learning/training_data.csv @@ -98,4 +98,54 @@ Irregular,Medium,No,High,High,Dark,Low,Yes, PLASTIC_AND_METAL Flat,High,Yes,Medium,Medium,Colorful,Low,Yes,PAPER Longitiudonal,Low,Yes,Medium,Medium,Light,Medium,Yes,MIXED Round,Medium,Yes,High,Low,Light,Medium,Yes,PLASTIC_AND_METAL -Longitiudonal,High,No,Medium,Low,Light,Low,Yes,PLASTIC_AND_METAL \ No newline at end of file +Longitiudonal,High,No,Medium,Low,Light,Low,Yes,PLASTIC_AND_METAL +Flat,Low,Yes,Low,High,Light,Medium,Yes,PLASTIC_AND_METAL +Irregular,High,No,Low,High,Light,Low,No,PAPER +Round,High,Yes,High,Medium,Colorful,Low,No,MIXED +Flat,Low,Yes,Medium,Low,Transparent,Low,Yes,PLASTIC_AND_METAL +Irregular,Low,No,Medium,Low,Dark,Low,No,PAPER +Round,Low,No,Medium,Medium,Light,Medium,Yes,PLASTIC_AND_METAL +Longitiudonal,High,Yes,Low,Low,Light,High,No,MIXED +Irregular,Medium,Yes,Low,Medium,Light,High,No,MIXED +Flat,High,Yes,High,Medium,Dark,Low,No,PLASTIC_AND_METAL +Round,Low,Yes,Medium,Medium,Light,Medium,No,PLASTIC_AND_METAL +Round,Low,No,Low,Medium,Dark,Medium,No,PAPER +Irregular,Low,Yes,High,Low,Transparent,Medium,No,PLASTIC_AND_METAL +Longitiudonal,High,Yes,High,High,Transparent,Low,Yes,PLASTIC_AND_METAL +Round,High,Yes,Low,Medium,Transparent,Low,Yes,MIXED +Longitiudonal,High,Yes,Medium,Medium,Colorful,Low,Yes,MIXED +Round,Low,Yes,Low,Medium,Dark,Low,Yes,MIXED +Flat,Medium,No,Medium,High,Colorful,Low,No,PAPER +Irregular,High,Yes,Low,High,Transparent,Medium,No,BIO +Round,Medium,Yes,High,Medium,Transparent,High,Yes,MIXED +Flat,Low,No,Medium,Medium,Colorful,High,No,GLASS +Longitiudonal,Medium,No,High,Medium,Colorful,Medium,No,GLASS +Longitiudonal,Low,Yes,Low,Medium,Colorful,Medium,Yes,PLASTIC_AND_METAL +Irregular,Medium,No,Medium,High,Transparent,Medium,No,PAPER +Longitiudonal,Medium,Yes,Low,Medium,Dark,High,Yes,MIXED +Round,Low,Yes,High,High,Dark,Medium,No,MIXED +Longitiudonal,High,No,Low,High,Colorful,High,No,MIXED +Flat,Low,Yes,Medium,Medium,Transparent,High,Yes,PLASTIC_AND_METAL +Irregular,High,Yes,High,Medium,Light,High,No,PAPER +Longitiudonal,Low,Yes,High,Medium,Colorful,Low,Yes,PAPER +Round,Medium,No,High,Low,Colorful,High,No,PLASTIC_AND_METAL +Longitiudonal,High,No,Medium,Low,Dark,Medium,No,MIXED +Round,High,Yes,Medium,Medium,Light,Medium,No,GLASS +Longitiudonal,Low,No,Low,Medium,Light,Medium,No,MIXED +Longitiudonal,High,Yes,High,Low,Transparent,High,Yes,PLASTIC_AND_METAL +Round,Medium,Yes,Low,High,Light,High,No,PLASTIC_AND_METAL +Round,High,Yes,High,High,Light,Medium,Yes,MIXED +Flat,Low,No,Medium,High,Transparent,High,Yes,PAPER +Irregular,Medium,Yes,Low,High,Dark,High,Yes,PLASTIC_AND_METAL +Longitiudonal,Low,No,Low,Low,Transparent,High,No,MIXED +Round,High,No,Medium,High,Light,High,Yes,PAPER +Irregular,Low,No,Low,High,Light,Medium,Yes,MIXED +Irregular,High,No,Low,High,Light,High,No,PAPER +Longitiudonal,High,No,Low,Medium,Transparent,Low,No,PLASTIC_AND_METAL +Irregular,Low,Yes,High,Low,Colorful,Medium,No,MIXED +Round,Low,No,High,Low,Colorful,High,No,BIO +Irregular,Medium,No,High,High,Colorful,Low,Yes,BIO +Irregular,High,No,Medium,High,Colorful,Low,Yes,PAPER +Irregular,Low,No,Medium,Low,Dark,Medium,Yes,PLASTIC_AND_METAL +Longitiudonal,Medium,Yes,Low,Low,Light,High,Yes,MIXED +Flat,High,No,High,Low,Transparent,Low,No,PLASTIC_AND_METAL