243 lines
7.7 KiB
Python
243 lines
7.7 KiB
Python
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
|