55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
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(1, 25)
|
|
self.space = random.randint(1, 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
|