diff --git a/models/dumpster.py b/models/dumpster.py index d88dc2f..8402a31 100644 --- a/models/dumpster.py +++ b/models/dumpster.py @@ -1,69 +1,109 @@ import numpy as np -import pygame +import random +import os + class trash(): def __init__(self): self.id = 0 self.plastic = 0 self.paper = 0 - self.mixed = 0 - self.organic = 0 + self.metal = 0 + self.cardboard = 0 self.glass = 0 self.color = 0 self.xy = 0 self.value = 1 self.numberx = 0 self.numbery = 0 + self.trashPaths = [] + self.trash = 0 + self.kind = "" def giveID(self,id): self.id = id + #puts garbage locations into a list + def createTrash(self, x, type): + trashPaths=[] + for i in range(x): + kind = random.choice(type) + path = "Garbage classification\\testset\\" + kind + file = random.choice(os.listdir(path)) + path = "Garbage classification\\testset\\" + kind + "\\" + file + trashPaths.append(path) + return trashPaths def throwAway(self): trash = np.random.randint(1,6) if trash == 1: - self.plastic += 1 + #plastic YELLOW + tra = self.createTrash(random.randint(1,3), ["plastic"]) + sh = self.createTrash((random.randint(0,2)), ["cardboard", "glass", "metal", "paper", "plastic"]) + self.trash = tra + sh + self.kind = "plastic" + self.plastic += len(self.trash) return True elif trash == 2: - self.paper += 1 + #paper GREEN + tra = self.createTrash(random.randint(1,3), ["paper"]) + sh = self.createTrash((random.randint(0,2)), ["cardboard", "glass", "metal", "paper", "plastic"]) + self.trash = tra + sh + self.kind = "paper" + self.paper += len(self.trash) return True elif trash == 3: - self.glass += 1 + #glass RED + tra = self.createTrash(random.randint(1,3), ["glass"]) + sh = self.createTrash((random.randint(0,2)), ["cardboard", "glass", "metal", "paper", "plastic"]) + self.trash = tra + sh + self.kind = "glass" + self.glass += len(self.trash) return True elif trash == 4: - self.organic += 1 + #cardboard BLUE + tra = self.createTrash(random.randint(1,3), ["cardboard"]) + sh = self.createTrash((random.randint(0,2)), ["cardboard", "glass", "metal", "paper", "plastic"]) + self.trash = tra + sh + self.kind = "cardboard" + self.cardboard += len(self.trash) return True elif trash == 5: - self.mixed += 1 + #metal BLACK + tra = self.createTrash(random.randint(1,3), ["metal"]) + sh = self.createTrash((random.randint(0,2)), ["cardboard", "glass", "metal", "paper", "plastic"]) + self.trash = tra + sh + self.kind = "metal" + self.metal += len(self.trash) return True else: return False - def empty(self): - if self.plastic > 0: + def empty(self, i): + + if i == "plastic": self.plastic -= 1 - print("taking plastic") + #print("taking plastic") return True - elif self.paper > 0: + elif i == "paper": self.paper -= 1 - print("taking paper") + #print("taking paper") return True - elif self.glass > 0: + elif i == "glass": self.glass -= 1 - print("taking glass") + #print("taking glass") return True - elif self.organic > 0: - self.organic -= 1 - print("taking organic") + elif i == "cardboard": + self.cardboard -= 1 + #print("taking cardboard") return True - elif self.mixed > 0: - self.mixed -= 1 - print("taking mixed") + elif i == "metal": + self.metal -= 1 + #print("taking metal") return True else: - print("can't take trash from this dumpster") return False def colour(self): @@ -76,10 +116,10 @@ class trash(): elif self.glass > 0: self.color = 3 return True - elif self.organic > 0: + elif self.cardboard > 0: self.color = 4 return True - elif self.mixed > 0: + elif self.metal > 0: self.color = 5 return True else: