diff --git a/AI_brain/image_recognition.py b/AI_brain/image_recognition.py new file mode 100644 index 0000000..58ba878 --- /dev/null +++ b/AI_brain/image_recognition.py @@ -0,0 +1,42 @@ +import os +import numpy as np +import tensorflow as tf +from tensorflow import keras +import cv2 +import random + + +class VacuumRecognizer: + model = keras.models.load_model("D:/Image_dataset/model.h5") + + def recognize(self, image_path) -> str: + class_names = ['Banana', 'Cat', 'Earings', 'Plant'] + + img = cv2.imread(image_path, flags=cv2.IMREAD_GRAYSCALE) + # print(img.shape) + cv2.imshow("lala", img) + cv2.waitKey(0) + img = (np.expand_dims(img, 0)) + + predictions = self.model.predict(img)[0].tolist() + + print(class_names) + print(predictions) + print(max(predictions)) + print(predictions.index(max(predictions))) + + return class_names[predictions.index(max(predictions))] + + +image_paths = [] +image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Banana/') +image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Cat/') +image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Earings/') +image_paths.append('D:/Image_dataset/Image_datasetJPGnewBnW/Image_datasetJPGnewBnW/test/Plant/') +uio = VacuumRecognizer() + + +for image_path in image_paths: + dirs = os.listdir(image_path) + for i in range(10): + print(uio.recognize(image_path + dirs[random.randint(0, len(dirs)-1)])) \ No newline at end of file diff --git a/main.py b/main.py index 2b75c16..00dee9b 100644 --- a/main.py +++ b/main.py @@ -28,7 +28,7 @@ class Main: tiles_x = 10 tiles_y = 10 - self.renderer = Renderer(800, 800, tiles_x, tiles_y) + self.renderer = Renderer(600, 600, tiles_x, tiles_y) self.world = generate_world(tiles_x, tiles_y)