47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
|
import os
|
||
|
|
import random
|
||
|
|
from pathlib import Path
|
||
|
|
import numpy as np
|
||
|
|
import tensorflow as tf
|
||
|
|
from tensorflow import keras
|
||
|
|
|
||
|
|
|
||
|
|
class Predictor:
|
||
|
|
def __init__(self):
|
||
|
|
# Load the trained model
|
||
|
|
self.model = keras.models.load_model('Network/trained_model.h5')
|
||
|
|
|
||
|
|
# Load the class names
|
||
|
|
self.class_names = ['table', 'done', '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)
|
||
|
|
predicted_class_index = np.argmax(predictions[0])
|
||
|
|
predicted_class = self.class_names[predicted_class_index]
|
||
|
|
|
||
|
|
print(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
|