bob-v1-test / model_utils.py
ivanmichael's picture
Create model_utils.py
c0c9b22 verified
# model_utils.py
import tensorflow as tf
from PIL import Image
import numpy as np
class ImageClassifier:
def __init__(self, model_path):
self.model = tf.keras.models.load_model(model_path)
# Update these based on your model's requirements
self.input_size = (224, 224) # Example size, change as needed
self.class_names = ['class1', 'class2', 'class3'] # Replace with your class names
def preprocess_image(self, image):
"""Preprocess the image for your model"""
image = image.resize(self.input_size)
image_array = np.array(image)
image_array = image_array / 255.0 # Normalize if your model expects this
image_array = np.expand_dims(image_array, axis=0)
return image_array
def predict(self, image):
"""Make a prediction on the image"""
processed_image = self.preprocess_image(image)
predictions = self.model.predict(processed_image)
predicted_class = np.argmax(predictions[0])
confidence = np.max(predictions[0])
return {
'class': self.class_names[predicted_class],
'confidence': float(confidence),
'all_predictions': predictions.tolist()
}