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:
commit
82b964ac2a
42
AI_brain/image_recognition.py
Normal file
42
AI_brain/image_recognition.py
Normal 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)]))
|
Loading…
Reference in New Issue
Block a user