implementation for 2nd individual project

This commit is contained in:
Kamila Bobkowska 2020-05-07 16:06:30 +00:00
parent fc6a719351
commit edfdbc255f

View File

@ -1,105 +1,222 @@
import pygame import pygame
import random import random
from models.dumpster import trash from models.dumpster import trash
class GarbageTruck():
def __init__(self): #libraries used in the CNN part
self.coor = [0,0] import os
self.capacity = 5 import numpy as np
self.plastic = 0 import random
self.paper = 0 import shutil
self.mixed = 0 from keras.models import Sequential
self.organic = 0 from keras.layers import Conv2D, Flatten, MaxPooling2D, Dense
self.glass = 0 from keras.preprocessing import image
self.trash=trash() from keras.preprocessing.image import ImageDataGenerator
self.dir = 'R' import matplotlib.pyplot as plt
self.picture=pygame.image.load('./images/truckR.png')
def start(self): class GarbageTruck():
self.coor = [random.randint(0,19),random.randint(0,19)] def __init__(self):
return self.coor self.coor = [0,0]
self.capacity = 5
def amIFull(self): self.plastic = 0
#check id garbage truck full self.paper = 0
if ( self.plastic==self.capacity and self.paper==self.capacity and self.metal = 0
self.glass==self.capacity and self.mixed==self.capacity and self.cardboard = 0
self.organic==self.capacity): self.glass = 0
return True self.trash=trash()
#otherwise see if a category is full self.dir = 'R'
elif self.plastic!=self.capacity: self.picture=pygame.image.load('./images/truckR.png')
return False
elif self.paper!=self.capacity: def start(self):
return False self.coor = [random.randint(0,19),random.randint(0,19)]
elif self.glass!=self.capacity: return self.coor
return False
elif self.mixed!=self.capacity: #check id garbage truck full
return False def amIFull(self):
elif self.organic!=self.capacity: if ( self.plastic==self.capacity and self.paper==self.capacity and
return False self.glass==self.capacity and self.metal==self.capacity and
self.cardboard==self.capacity):
def move(self): print("DUMPSTER TRUCK FULL")
if(self.dir=="R"): return False
self.coor[0]+=1 else:
elif(self.dir=="L"): return True
self.coor[0]-=1
elif(self.dir=="U"): # used to move in A*
self.coor[1]-=1 def move(self):
elif(self.dir=="D"): if(self.dir=="R"):
self.coor[1]+=1 self.coor[0]+=1
elif(self.dir=="L"):
def turn(self,goal): self.coor[0]-=1
if(goal[0]-self.coor[0]==1): elif(self.dir=="U"):
self.dir='R' self.coor[1]-=1
self.picture = pygame.image.load('./images/truckR.png') elif(self.dir=="D"):
elif(goal[0]-self.coor[0]==-1): self.coor[1]+=1
self.dir='L'
self.picture = pygame.image.load('./images/truckL.png') # used to move in A*
elif(goal[1]-self.coor[1]==-1): def turn(self,goal):
self.dir='U' if(goal[0]-self.coor[0]==1):
self.picture = pygame.image.load('./images/truckU.png') self.dir='R'
elif(goal[1]-self.coor[1]==1): self.picture = pygame.image.load('./images/truckR.png')
self.dir='D' elif(goal[0]-self.coor[0]==-1):
self.picture = pygame.image.load('./images/truckD.png') self.dir='L'
return dir self.picture = pygame.image.load('./images/truckL.png')
elif(goal[1]-self.coor[1]==-1):
''' #old move self.dir='U'
def move(self, coor, moving, size): self.picture = pygame.image.load('./images/truckU.png')
if(moving == 0 and coor[0] > 0): elif(goal[1]-self.coor[1]==1):
coor[0] = coor[0] - 1 self.dir='D'
return(coor) self.picture = pygame.image.load('./images/truckD.png')
elif(moving == 1 and coor[1] > 0): return dir
coor[1] = coor[1] - 1
return(coor) # main function for an attempt to collect garbage from a dumpster
elif(moving == 2 and coor[0] < size - 1): def collectingTrash(self, trash, id):
coor[0] = coor[0] + 1 if self.amIFull():
return(coor) inthedumpster = self.check(trash, id)
elif(moving == 3 and coor[1] < size - 1): self.takingTrash(trash, id, inthedumpster)
coor[1] = coor[1] + 1
return(coor) #trying to put a specific piece of trash if possible
else: def collect(self, trash, id, i):
return(coor)''' if i == "plastic":
if self.plastic < self.capacity:
def collectingTrash(self, trash, id): #print("Plastic put into garbage truck")
#1st if recognizes the dumpster kind self.plastic+=1
if trash[id].color==1: return True
#2nd if check if there is space in the garbage truck and if there is still trash else:
if self.plastic<self.capacity and trash[id].plastic > 0: #print("No more space in garbage truck")
self.plastic+=1 return False
return True
elif trash[id].color==2: elif i == "paper":
if self.paper<self.capacity and trash[id].paper > 0: if self.paper < self.capacity:
self.paper+=1 #print("Paper put into garbage truck")
return True self.paper+=1
elif trash[id].color==3: return True
if self.glass<self.capacity and trash[id].glass > 0: else:
self.glass+=1 #print("No more space in garbage truck")
return True return False
elif trash[id].color==4:
if self.organic<self.capacity and trash[id].organic > 0: elif i == "glass":
self.organic+=1 if self.glass < self.capacity:
return True #print("Glass put into garbage truck")
elif trash[id].color==5: self.glass+=1
if self.mixed<self.capacity and trash[id].mixed > 0: return True
self.mixed+=1 else:
return True #print("No more space in garbage truck")
else: return False
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 = "binary_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_4.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