AI2020_Project/models/Garbagetruck.py

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