AI2020_Project/models/Garbagetruck.py

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

    #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