Machine_learning_2023/AI_brain/image_recognition.py
2023-06-01 19:35:54 +02:00

42 lines
1.3 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("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)]))