AI2020_Project/models/dumpster.py

import numpy as np
import random
import os


class trash():
    def __init__(self):
        self.id = 0
        self.plastic = 0
        self.paper = 0
        self.metal = 0
        self.cardboard = 0
        self.glass = 0
        self.color = 0
        self.xy = 0
        self.value = 1
        self.numberx = 0
        self.numbery = 0
        self.trashPaths = []
        self.trash = 0
        self.kind = ""

    def giveID(self,id):
        self.id = id

    #puts garbage locations into a list
    def createTrash(self, x, type):
        trashPaths=[]
        for i in range(x):
            kind = random.choice(type)
            path = "Garbage classification\\testset\\" + kind
            file = random.choice(os.listdir(path))
            path = "Garbage classification\\testset\\" + kind + "\\" + file
            trashPaths.append(path)
        return trashPaths


    def throwAway(self):
        trash = np.random.randint(1,6)
        if trash == 1:
            #plastic YELLOW
            tra = self.createTrash(random.randint(1,3), ["plastic"])
            sh = self.createTrash((random.randint(0,2)), ["cardboard", "glass", "metal", "paper", "plastic"])
            self.trash = tra + sh
            self.kind = "plastic"
            self.plastic += len(self.trash)
            return True
        elif trash == 2:
            #paper GREEN
            tra = self.createTrash(random.randint(1,3), ["paper"])
            sh = self.createTrash((random.randint(0,2)), ["cardboard", "glass", "metal", "paper", "plastic"])
            self.trash = tra + sh
            self.kind = "paper"
            self.paper += len(self.trash)
            return True
        elif trash == 3:
            #glass RED
            tra = self.createTrash(random.randint(1,3), ["glass"])
            sh = self.createTrash((random.randint(0,2)), ["cardboard", "glass", "metal", "paper", "plastic"])
            self.trash = tra + sh
            self.kind = "glass"
            self.glass += len(self.trash)
            return True
        elif trash == 4:
            #cardboard BLUE
            tra = self.createTrash(random.randint(1,3), ["cardboard"])
            sh = self.createTrash((random.randint(0,2)), ["cardboard", "glass", "metal", "paper", "plastic"])
            self.trash = tra + sh
            self.kind = "cardboard"
            self.cardboard += len(self.trash)
            return True
        elif trash == 5:
            #metal BLACK
            tra = self.createTrash(random.randint(1,3), ["metal"])
            sh = self.createTrash((random.randint(0,2)), ["cardboard", "glass", "metal", "paper", "plastic"])
            self.trash = tra + sh
            self.kind = "metal"
            self.metal += len(self.trash)
            return True
        else:
            return False


    def empty(self, i):

        if i == "plastic":
            self.plastic -= 1
            #print("taking plastic")
            return True
        elif i == "paper":
            self.paper -= 1
            #print("taking paper")
            return True
        elif i == "glass":
            self.glass -= 1
            #print("taking glass")
            return True
        elif i == "cardboard":
            self.cardboard -= 1
            #print("taking cardboard")
            return True
        elif i == "metal":
            self.metal -= 1
            #print("taking metal")
            return True
        else:
            return False

    def colour(self):
        if self.plastic > 0:
            self.color = 1
            return True
        elif self.paper > 0:
            self.color = 2
            return True
        elif self.glass > 0:
            self.color = 3
            return True
        elif self.cardboard > 0:
            self.color = 4
            return True
        elif self.metal > 0:
            self.color = 5
            return True
        else:
            return False