automatyczny_kelner/Network/Predictor.py

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