Merge pull request 'image_recognition' (#26) from image_recognition into main

Reviewed-on: #26
Reviewed-by: Tim Barvenov <timbar@st.amu.edu.pl>
This commit is contained in:
Nastassia Zhuravel 2023-06-01 19:38:09 +02:00
commit 82b964ac2a
2 changed files with 43 additions and 1 deletions

View File

@ -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)]))

View File

@ -28,7 +28,7 @@ class Main:
tiles_x = 10 tiles_x = 10
tiles_y = 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) self.world = generate_world(tiles_x, tiles_y)