Machine_learning_2023/AI_brain/image_recognition.py
2023-06-14 19:59:23 +02:00

82 lines
3.4 KiB
Python

import os
import numpy as np
import tensorflow as tf
from tensorflow import keras
import cv2
import random
# You can download model from https://uam-my.sharepoint.com/:f:/g/personal/pavbia_st_amu_edu_pl/EmBHjnETuk5LiCZS6xk7AnIBNsnffR3Sygf8EX2bhR1w4A
# Change the path to model + to datasets (string 12 + strings 35,41,47,53)
class VacuumRecognizer:
model = keras.models.load_model("AI_brain\model.h5") # Neuron model path
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()
# More information about neuron decision
# 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))]
def get_random_dir(self, type) -> str:
if type == "Plant":
plant_image_paths = "C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Plant" # Plant dataset path
plant_dirs = os.listdir(plant_image_paths)
full_path = (
plant_image_paths
+ "\\"
+ plant_dirs[random.randint(0, len(plant_dirs) - 1)]
)
print(full_path)
return full_path
elif type == "Earings":
earnings_image_paths = "C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Earings" # Earings dataset path
earning_dirs = os.listdir(earnings_image_paths)
full_path = (
earnings_image_paths
+ "\\"
+ earning_dirs[random.randint(0, len(earning_dirs) - 1)]
)
print(full_path)
return full_path
elif type == "Banana":
banana_image_paths = "C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Banana" # Banana dataset path
banana_dirs = os.listdir(banana_image_paths)
full_path = (
banana_image_paths
+ "\\"
+ banana_dirs[random.randint(0, len(banana_dirs) - 1)]
)
print(full_path)
return full_path
elif type == "Cat":
cat_image_paths = "C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Cat" # Cat dataset path
cat_dir = os.listdir(cat_image_paths)
# 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)]))"""