drzewo decyzyjne

main.py

@@ -4,6 +4,7 @@ from queue import PriorityQueue
 from path_algorithms.a_star import a_star
 # from path_algorithms.bfs import bfs
 from rubbish import *
+from tree import evaluate_values, trash_selection
 from truck import Truck
 from surface import *
 
@@ -49,6 +50,7 @@ for i in range(15):
             rubbish_list.append(Rubbish(screen, j * 60, i * 60))
 
 path = []
+run = 1
 while True:
     pygame.time.delay(500)
 
@@ -58,6 +60,13 @@ while True:
         i.draw_rubbish()
     truck.draw_truck()
 
+    if run == 1:
+        # func(rubbish_list[0])
+        data = rubbish_list[0].data_for_decision_tree()
+        test = trash_selection(evaluate_values(data))
+        print(test)
+        run = 0
+
     if rubbish_list and not path:
         start = (truck.y / 60, truck.x / 60)
         direction = truck.direction

rubbish.py

@@ -13,49 +13,52 @@ class Rubbish:
         self.image = pygame.image.load('images/garbage.png')
         self.surface_rect = self.image.get_rect()
         self.surface_rect.center = (self.x + 30, self.y + 30)
-        self.weight = random.randint(0, 10)
+        self.weight = random.choice([10, 15, 20, 30, 40])
-        self.density = random.randint(0, 10)
+        self.density = random.choice([True, False])
-        self.fragility = random.randint(0, 10)
+        self.fragility = random.randint(1, 10)
-        self.dirty = random.randint(0, 10)
+        self.material = random.choice(["plastic", "wood", "metal", "glass", "paper"])
-        self.size = random.randint(0, 10)
+        self.size = random.choice(["little", "medium", "huge", "large"])
-        self.degradability = random.randint(0, 10)
+        self.degradability = random.randint(1, 10)
-        self.renewability = random.randint(0, 10)
+        self.renewability = random.randint(1, 10)
 
     def draw_rubbish(self):
         self.screen.blit(self.image, self.surface_rect)
 
+    def data_for_decision_tree(self):
+        return [self.weight, self.density, self.fragility, self.material, self.size, self.degradability, self.renewability]
 
-class Paper_waste(Rubbish):
+
+class PaperWaste(Rubbish):
 
     def __init__(self, screen, x, y):
         super().__init__(screen, x, y)
 
 
-class Organic_waste(Rubbish):
+class OrganicWaste(Rubbish):
 
     def __init__(self, screen, x, y):
         super().__init__(screen, x, y)
 
 
-class Glass_waste(Rubbish):
+class GlassWaste(Rubbish):
 
     def __init__(self, screen, x, y):
         super().__init__(screen, x, y)
 
 
-class Plastic_waste(Rubbish):
+class PlasticWaste(Rubbish):
 
     def __init__(self, screen, x, y):
         super().__init__(screen, x, y)
 
 
-class E_waste(Rubbish):
+class EWaste(Rubbish):
 
     def __init__(self, screen, x, y):
         super().__init__(screen, x, y)
 
 
-class Metal_waste(Rubbish):
+class MetalWaste(Rubbish):
 
     def __init__(self, screen, x, y):
         super().__init__(screen, x, y)

tree.py

@@ -0,0 +1,151 @@
+import random
+
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from sklearn.tree import DecisionTreeClassifier
+
+
+def evaluate_values(values):
+    data = []
+    if values[0] == 10:
+        data.append(10)
+    elif values[0] == 15:
+        data.append(15)
+    elif values[0] == 20:
+        data.append(20)
+    elif values[0] == 30:
+        data.append(30)
+    elif values[0] == 40:
+        data.append(40)
+    else:
+        data.append(random.choice([10, 15, 20, 30, 40]))
+
+    if values[1] is True:
+        data.append(1)
+    elif values[1] is False:
+        data.append(0)
+    else:
+        data.append(random.choice([1, 0]))
+
+    if values[2] == 1:
+        data.append(1)
+    elif values[2] == 2:
+        data.append(2)
+    elif values[2] == 3:
+        data.append(3)
+    elif values[2] == 4:
+        data.append(4)
+    elif values[2] == 5:
+        data.append(5)
+    elif values[2] == 6:
+        data.append(6)
+    elif values[2] == 7:
+        data.append(7)
+    elif values[2] == 8:
+        data.append(8)
+    elif values[2] == 9:
+        data.append(9)
+    elif values[2] == 10:
+        data.append(10)
+    else:
+        data.append(random.choice([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]))
+
+    if values[3] == 'plastic':
+        data.append(5)
+    elif values[3] == 'wood':
+        data.append(2)
+    elif values[3] == 'metal':
+        data.append(4)
+    elif values[3] == 'glass':
+        data.append(3)
+    elif values[3] == 'paper':
+        data.append(1)
+    else:
+        data.append(random.choice([1, 2, 3, 4, 5]))
+
+    if values[4] == 'little':
+        data.append(1)
+    elif values[4] == 'medium':
+        data.append(2)
+    elif values[4] == 'huge':
+        data.append(3)
+    elif values[4] == 'large':
+        data.append(4)
+    else:
+        data.append(random.choice([1, 2, 3, 4]))
+
+    if values[5] == 1:
+        data.append(1)
+    elif values[5] == 2:
+        data.append(2)
+    elif values[5] == 3:
+        data.append(3)
+    elif values[5] == 4:
+        data.append(4)
+    elif values[5] == 5:
+        data.append(5)
+    elif values[5] == 6:
+        data.append(6)
+    elif values[5] == 7:
+        data.append(7)
+    elif values[5] == 8:
+        data.append(8)
+    elif values[5] == 9:
+        data.append(9)
+    elif values[5] == 10:
+        data.append(10)
+    else:
+        data.append(random.choice([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]))
+
+    if values[6] == 1:
+        data.append(1)
+    elif values[6] == 2:
+        data.append(2)
+    elif values[6] == 3:
+        data.append(3)
+    elif values[6] == 4:
+        data.append(4)
+    elif values[6] == 5:
+        data.append(5)
+    elif values[6] == 6:
+        data.append(6)
+    elif values[6] == 7:
+        data.append(7)
+    elif values[6] == 8:
+        data.append(8)
+    elif values[6] == 9:
+        data.append(9)
+    elif values[6] == 10:
+        data.append(10)
+    else:
+        data.append(random.choice([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]))
+    return data
+
+
+def trash_selection(prefer):
+    df = pd.read_excel('data.xlsx', sheet_name='list1')
+    # print(df)
+
+    d = {'paper': 1, 'wood': 2, 'glass': 3, 'metal': 4, 'plastic': 5}
+    df['material'] = df['material'].map(d)
+
+    d = {'little': 1, 'medium': 2, 'huge': 3, 'large': 4}
+    df['size'] = df['size'].map(d)
+
+    d = {'leave': 0, 'pick up': 1}
+    df['what to do'] = df['what to do'].map(d)
+
+    features = ['weight', 'density', 'fragility', 'material', 'size', 'degradability', 'renewability']
+    x = df[features]
+    y = df['what to do']
+    x_train, x_test, y_train, y_test = train_test_split(x, y)
+
+    clf = DecisionTreeClassifier(criterion='entropy')
+
+    # model = clf.fit(X, y)
+    # text_representation = tree.export_text(clf)
+    # print(text_representation)
+
+    clf = clf.fit(x_train, y_train)
+
+    return clf.predict([prefer])