import os import random import logging import numpy as np import tensorflow as tf from tensorflow import keras from termcolor import colored class Predictor: def __init__(self): # Turn off interactive logging tf.get_logger().setLevel(logging.ERROR) # Load the trained model self.model = keras.models.load_model('Network/trained_model.h5') # Load the class names self.class_names = ['table', 'table', 'order'] # Path to the folder containing test images self.test_images_folder = 'Network/Testing/' def predict(self, image_path): # Load and preprocess the test image test_image = keras.preprocessing.image.load_img( image_path, target_size=(100, 100)) test_image = keras.preprocessing.image.img_to_array(test_image) test_image = np.expand_dims(test_image, axis=0) test_image = test_image / 255.0 # Normalize the image # Reshape the image array to (1, height, width, channels) test_image = np.reshape(test_image, (1, 100, 100, 3)) # Make predictions predictions = self.model.predict(test_image, verbose=None) predicted_class_index = np.argmax(predictions[0]) predicted_class = self.class_names[predicted_class_index] print(colored("Predicted class: ", "yellow")+f"{predicted_class}") return predicted_class def random_path_img(self) -> str: folder_name = random.choice(os.listdir(self.test_images_folder)) folder_path = os.path.join(self.test_images_folder, folder_name) filename = "" while not (filename.endswith('.jpg') or filename.endswith('.jpeg')): filename = random.choice(os.listdir(folder_path)) image_path = os.path.join(folder_path, filename) return image_path