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()
-            
+
+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')

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
+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