43 lines
1.5 KiB
Python
43 lines
1.5 KiB
Python
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('AI_brain\model.h5')
|
|
|
|
def recognize(self, image_path) -> str:
|
|
class_names = ['Banana', 'Cat', 'Earings', 'Plant']
|
|
|
|
img = cv2.imread(image_path, flags=cv2.IMREAD_GRAYSCALE)
|
|
cv2.waitKey(0)
|
|
img = (np.expand_dims(img, 0))
|
|
|
|
predictions = self.model.predict(img)[0].tolist()
|
|
|
|
# print(img.shape)
|
|
# cv2.imshow("test_show", img)
|
|
# print(class_names)
|
|
# print(predictions)
|
|
# print(max(predictions))
|
|
# print(predictions.index(max(predictions)))
|
|
|
|
return class_names[predictions.index(max(predictions))]
|
|
|
|
|
|
|
|
#For testing the neuron model
|
|
'''image_paths = []
|
|
image_paths.append('C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Banana')
|
|
image_paths.append('C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Cat')
|
|
image_paths.append('C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Earings')
|
|
image_paths.append('C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Plant')
|
|
uio = VacuumRecognizer()
|
|
|
|
for image_path in image_paths:
|
|
dirs = os.listdir(image_path)
|
|
for i in range(3):
|
|
print(uio.recognize(image_path + '\\' + dirs[random.randint(0, len(dirs)-1)]))''' |