215 lines
6.1 KiB
Python
215 lines
6.1 KiB
Python
import pygame as pg
|
|
import numpy as np
|
|
import random as rd
|
|
from os import listdir
|
|
from os.path import isfile, join
|
|
# from Logic.TrashRecognition.ImageClassification import classify
|
|
|
|
# MODULE LEVEL VARIABLES
|
|
recognized_trash = {
|
|
|
|
}
|
|
# trash_files = classify()
|
|
########################
|
|
|
|
|
|
class Grid:
|
|
|
|
def __init__(self, cols: int, rows: int):
|
|
self.table = [[Node(row, col)
|
|
for col in range(cols)]
|
|
for row in range(rows)]
|
|
self.cols = cols
|
|
self.rows = rows
|
|
|
|
def draw_map(self, screen: "PyGame screen"):
|
|
"""Draws whole map"""
|
|
screen.fill(Node.BLACK)
|
|
for row in self.table:
|
|
for node in row:
|
|
node.draw(screen)
|
|
|
|
def change_field(self, row: int, col: int, f_type: int):
|
|
self.table[row][col].change_field_type(f_type)
|
|
|
|
def generate_trash(self,row: int, col: int):
|
|
self.table[row][col].generate_trash()
|
|
|
|
def draw_node(self, screen, row: int, col: int):
|
|
self.table[row][col].draw(screen)
|
|
|
|
def get_trash_possition(self, trash: int):
|
|
trash_possition = []
|
|
for row in self.table:
|
|
for node in row:
|
|
if node.field_type == trash and node.house.empty == False:
|
|
trash_possition.append((node.row,node.col))
|
|
return trash_possition
|
|
|
|
def garbage_to_collect(self):
|
|
garbage = []
|
|
field_types = []
|
|
d = rd.randint(1,7)
|
|
day = self.table[0][0].house.get_day_of_week(d)
|
|
for index, x in enumerate(day):
|
|
if index != 0 and index != 1:
|
|
garbage.append(x)
|
|
print("\n Today is:", day[1], ", garbage to collect: ",end = '')
|
|
for x in garbage:
|
|
print(x[2],", ",end='')
|
|
field_types.append(x[0])
|
|
return field_types
|
|
|
|
|
|
|
|
|
|
class House:
|
|
# define some basic colors
|
|
BLACK = (0, 0, 0)
|
|
WHITE = (255, 255, 255)
|
|
GREEN = (0, 255, 0)
|
|
SKYBLUE = (0,191,255)
|
|
RED = (255, 0, 0)
|
|
BLUE = (0, 0, 255)
|
|
ORANGE = (255, 165, 0)
|
|
PINK = (199,21,133)
|
|
GREY = (192,192,192)
|
|
|
|
#define trash
|
|
trash_dict = {
|
|
"paper": (5, WHITE, "paper"),
|
|
"glass": (6, SKYBLUE, "glass"),
|
|
"metal": (7, GREY, "metal"),
|
|
"plastic": (8, ORANGE, "plastic")
|
|
}
|
|
#define days of the week
|
|
MONDAY = (1, "Monday", trash_dict["paper"], trash_dict["metal"])
|
|
TUESDAY = (2, "Tuesday", trash_dict["glass"])
|
|
WEDNESDAY=(3, "Wednesday", trash_dict["plastic"], trash_dict["metal"])
|
|
THURSDAY = (4, "Thursday", trash_dict["glass"])
|
|
FRIDAY = (5, "Friday", trash_dict["paper"], trash_dict["metal"])
|
|
SATURDAY = (6, "Saturday", trash_dict["plastic"])
|
|
SUNDAY = (7, "Sunday", trash_dict["metal"])
|
|
DAYS = [MONDAY, TUESDAY, WEDNESDAY, THURSDAY, FRIDAY, SATURDAY, SUNDAY]
|
|
|
|
|
|
def __init__(self):
|
|
self.empty = True
|
|
self.trash = None
|
|
self.trash_file = None
|
|
|
|
|
|
def find_trash_file(self):
|
|
from os.path import sep # culture and os invariant separator
|
|
|
|
file_names = [(join(f"Images{sep}TestImages", f))
|
|
for f in listdir(f"Images{sep}TestImages")
|
|
if isfile(join(f"Images{sep}TestImages", f))]
|
|
|
|
file_name = file_names[rd.randint(0,len(file_names)) - 1]
|
|
|
|
from Logic.TrashRecognition.ImageClassification import classify_file
|
|
|
|
if file_name in recognized_trash:
|
|
rt = recognized_trash[file_name]
|
|
return (file_name, rt[0], rt[1])
|
|
else:
|
|
classification = classify_file(file_dir=file_name)
|
|
recognized_trash[file_name] = (classification[1], classification[2])
|
|
return classification
|
|
|
|
def generate_trash(self):
|
|
self.empty = False
|
|
|
|
classification = self.find_trash_file()
|
|
self.trash = self.trash_dict[classification[1]]
|
|
self.trash_file = classification[0]
|
|
|
|
def get_day_of_week(self, d: int):
|
|
for day in self.DAYS:
|
|
if day[0] == d:
|
|
return day
|
|
|
|
|
|
|
|
class Node:
|
|
# define rectangles dimensions
|
|
r_width = 20
|
|
r_height = 20
|
|
r_margin = 5
|
|
|
|
# define some basic colors
|
|
# TODO: change to Enum
|
|
BLACK = (0, 0, 0)
|
|
WHITE = (255, 255, 255)
|
|
GREEN = (0, 255, 0)
|
|
SKYBLUE = (0,191,255)
|
|
RED = (255, 0, 0)
|
|
BLUE = (0, 0, 255)
|
|
ORANGE = (255, 165, 0)
|
|
PINK = (199,21,133)
|
|
GREY = (192,192,192)
|
|
|
|
|
|
def __init__(self, row: int, col: int,
|
|
field_type: int = 0):
|
|
self.row = row
|
|
self.col = col
|
|
self.field_type = field_type
|
|
self.house = House()
|
|
|
|
|
|
|
|
def draw(self, screen):
|
|
color = self.get_field_color()
|
|
col = self.col
|
|
row = self.row
|
|
width = self.r_width
|
|
height = self.r_height
|
|
margin = self.r_margin
|
|
# rect -> (left, top, width, height)
|
|
# draw.rect(surface, color, rect, margin)
|
|
pg.draw.rect(screen, color,
|
|
((col * (width + margin)) + margin,
|
|
(row * (height + margin)) + margin,
|
|
width, height))
|
|
pg.display.flip() # refresh screen
|
|
|
|
def generate_trash(self):
|
|
self.house.generate_trash()
|
|
self.field_type = self.house.trash[0]
|
|
|
|
|
|
def change_field_type(self, field_type: int):
|
|
self.field_type = field_type
|
|
|
|
def get_field_color(self) -> tuple:
|
|
"""Gets the color tuple of field"""
|
|
#base color
|
|
if self.field_type == 0:
|
|
return self.GREEN
|
|
#truck color
|
|
elif self.field_type == 1:
|
|
return self.RED
|
|
#end point
|
|
elif self.field_type == 2:
|
|
return self.BLUE
|
|
#obstacles color
|
|
elif self.field_type == 3:
|
|
return self.ORANGE
|
|
#path color
|
|
elif self.field_type == 4:
|
|
return self.PINK
|
|
#paper
|
|
elif self.field_type == 5:
|
|
return self.WHITE
|
|
#glass
|
|
elif self.field_type == 6:
|
|
return self.SKYBLUE
|
|
#metal
|
|
elif self.field_type == 7:
|
|
return self.GREY
|
|
#plastic
|
|
elif self.field_type == 8:
|
|
return self.ORANGE
|