SI_projekt_smieciarka/trash.py

from keras.models import Sequential, load_model
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Conv2D, MaxPooling2D
from keras.layers.normalization import BatchNormalization
from PIL import Image
import numpy as np
import os
import random
from house import is_house

IMAGE_SIZE = 256
IMAGE_DIRECTORY = './trash'
filenames = next(os.walk(IMAGE_DIRECTORY), (None, None, []))[2]

model = load_model("model.h5")

def draw_trash(files):
    random_trash = []
    random_trash_names = []
    random_trash_reshape = []
    for i in range(3):
        image_name = random.choice(files)
        image_path = os.path.join(IMAGE_DIRECTORY, image_name)
        img = Image.open(image_path)
        img = img.convert('L')
        img = img.resize((IMAGE_SIZE, IMAGE_SIZE), Image.ANTIALIAS)
        random_trash_names.append(image_name)
        random_trash.append(np.array(img))
    random_trash_reshape = np.array([i for i in random_trash]).reshape(-1, IMAGE_SIZE, IMAGE_SIZE, 1)
    return random_trash_reshape, random_trash_names


class Trash:
    def __init__(self, grid_w, grid_h, grid_size):
        self.grid_w = grid_w
        self.grid_h = grid_h
        self.size = grid_size
        self.content = draw_trash(filenames)[0]
        self.names = draw_trash(filenames)[1]
        self.mass = random.randint(0, 25)
        self.space = random.randint(0, 25)

    def new_pos(self, truck_pos, houses, multi):
        self.trash_content, self.trash_names = draw_trash(filenames)
        while True:
            self.pos = [random.randrange(0, self.grid_w, self.size),
                        random.randrange(0, self.grid_h, self.size)]
            if self.pos != truck_pos and not is_house(self.pos, houses) and self not in multi:
                break

    def __eq__(self, other):
        if isinstance(self, Trash):
            return self.pos == other.pos
        return False