diff --git a/data.xlsx b/data.xlsx new file mode 100644 index 0000000..2d3ade0 Binary files /dev/null and b/data.xlsx differ diff --git a/main.py b/main.py index 402e30e..2050ce8 100644 --- a/main.py +++ b/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 diff --git a/rubbish.py b/rubbish.py index 02edf73..aff1150 100644 --- a/rubbish.py +++ b/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.density = random.randint(0, 10) - self.fragility = random.randint(0, 10) - self.dirty = random.randint(0, 10) - self.size = random.randint(0, 10) - self.degradability = random.randint(0, 10) - self.renewability = random.randint(0, 10) + self.weight = random.choice([10, 15, 20, 30, 40]) + self.density = random.choice([True, False]) + self.fragility = random.randint(1, 10) + self.material = random.choice(["plastic", "wood", "metal", "glass", "paper"]) + self.size = random.choice(["little", "medium", "huge", "large"]) + self.degradability = random.randint(1, 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) diff --git a/tree.py b/tree.py new file mode 100644 index 0000000..dd0e063 --- /dev/null +++ b/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]) \ No newline at end of file