import pygame import random from models.dumpster import trash #libraries used in the CNN part import os import numpy as np import random import shutil from keras.models import Sequential from keras.layers import Conv2D, Flatten, MaxPooling2D, Dense from keras.preprocessing import image from keras.preprocessing.image import ImageDataGenerator import matplotlib.pyplot as plt class GarbageTruck(): def __init__(self): self.coor = [0,0] self.capacity = 5 self.plastic = 0 self.paper = 0 self.metal = 0 self.cardboard = 0 self.glass = 0 self.trash=trash() self.dir = 'R' self.picture=pygame.image.load('./images/truckR.png') def start(self): self.coor = [random.randint(0,19),random.randint(0,19)] return self.coor def empty(self, i): if i == "plastic": self.plastic -= 1 return True elif i == "paper": self.paper -= 1 return True elif i == "glass": self.glass -= 1 return True elif i == "cardboard": self.cardboard -= 1 return True elif i == "metal": self.metal -= 1 return True else: return False #check id garbage truck full def amIFull(self): if ( self.plastic==self.capacity and self.paper==self.capacity and self.glass==self.capacity and self.metal==self.capacity and self.cardboard==self.capacity): print("DUMPSTER TRUCK FULL") return False else: return True # used to move in A* def move(self): if(self.dir=="R"): self.coor[0]+=1 elif(self.dir=="L"): self.coor[0]-=1 elif(self.dir=="U"): self.coor[1]-=1 elif(self.dir=="D"): self.coor[1]+=1 # used to move in A* def turn(self,goal): if(goal[0]-self.coor[0]==1): self.dir='R' self.picture = pygame.image.load('./images/truckR.png') elif(goal[0]-self.coor[0]==-1): self.dir='L' self.picture = pygame.image.load('./images/truckL.png') elif(goal[1]-self.coor[1]==-1): self.dir='U' self.picture = pygame.image.load('./images/truckU.png') elif(goal[1]-self.coor[1]==1): self.dir='D' self.picture = pygame.image.load('./images/truckD.png') return dir # main function for an attempt to collect garbage from a dumpster def collectingTrash(self, trash, id): if self.amIFull(): inthedumpster = self.check(trash, id) self.takingTrash(trash, id, inthedumpster) #trying to put a specific piece of trash if possible def collect(self, trash, id, i): if i == "plastic": if self.plastic < self.capacity: #print("Plastic put into garbage truck") self.plastic+=1 return True else: #print("No more space in garbage truck") return False elif i == "paper": if self.paper < self.capacity: #print("Paper put into garbage truck") self.paper+=1 return True else: #print("No more space in garbage truck") return False elif i == "glass": if self.glass < self.capacity: #print("Glass put into garbage truck") self.glass+=1 return True else: #print("No more space in garbage truck") return False elif i == "cardboard": if self.cardboard < self.capacity: #print("Cardboard put into garbage truck") self.cardboard+=1 return True else: #print("No more space in garbage truck") return False elif i == "metal": if self.metal < self.capacity: #print("Metal put into garbage truck") self.metal+=1 return True else: #print("No more space in garbage truck") return False else: return False def cap(self): print("In garbage truck:") print("- Plastic:", self.plastic) print("- Paper:", self.paper) print("- Glass:", self.glass) print("- Cardboard:", self.cardboard) print("- Metal:", self.metal) def takingTrash(self, trash, id, insides): for i in insides: if trash[id].kind == i: print("This ", i," was correctly sorted") if self.collect(trash, id, i): trash[id].empty(i) else: print("This ", i , "shouldn't be here") if self.collect(trash, id, i): trash[id].empty(i) self.cap() return True def check(self, trash, id): classifier = Sequential() classifier.add(Conv2D(32, (3, 3), input_shape=(110, 110, 3), activation = "relu")) classifier.add(MaxPooling2D(pool_size = (2, 2))) classifier.add(Conv2D(64, (3, 3), activation = "relu")) classifier.add(MaxPooling2D(pool_size=(2, 2))) # this layer in ver 4 classifier.add(Conv2D(32, (3, 3), activation = "relu")) classifier.add(MaxPooling2D(pool_size=(2, 2))) # ----------------- classifier.add(Flatten()) classifier.add(Dense(activation = "relu", units = 64 )) classifier.add(Dense(activation = "softmax", units = 5)) classifier.compile(optimizer = "adam", loss = "categorical_crossentropy", metrics = ["accuracy"]) train_datagen = ImageDataGenerator( rescale=1./255, shear_range=0.1, zoom_range=0.1, width_shift_range=0.1, height_shift_range=0.1, horizontal_flip=True, vertical_flip=True, ) test_datagen = ImageDataGenerator( rescale=1./255, validation_split=0.1 ) train_generator = train_datagen.flow_from_directory( "Garbage classification\\trainset", target_size=(110, 110), batch_size=16, class_mode='categorical', #seed=0 ) test_generator = test_datagen.flow_from_directory( "Garbage classification\\testset", target_size=(110, 110), batch_size=16, class_mode='categorical', ) labels = (train_generator.class_indices) labels = dict((v,k) for k,v in labels.items()) test_x, test_y = test_generator.__getitem__(1) whatarewe = [] classifier.load_weights("model_ver_6.h5") ''' #optional for displaying the trash and what the CNN decided it is z=1 ''' for i in trash[id].trash: gz = image.load_img(i, target_size = (110,110)) ti = image.img_to_array(gz) ti=np.array(ti)/255.0 ti = np.expand_dims(ti, axis = 0) prediction = classifier.predict(ti) whatarewe.append(labels[np.argmax(prediction)]) ''' #optional for displaying the trash and what the CNN decided it is ax = plt.subplot(1, len(trash[id].trash), z) z+=1 plt.imshow(gz) ax.set_xticks([]) ax.set_yticks([]) plt.title("AI thinks:\n %s" % (labels[np.argmax(prediction)])) plt.suptitle("dumpster kind: %s" %(trash[id].kind) ) plt.show(block=False) plt.pause(3) plt.close() ''' return whatarewe