TireClassifier / app.py
matemoral's picture
Upload app.py
833c8c4 verified
Raw
History Blame Contribute Delete
2.34 kB
import tensorflow as tf
import numpy as np
import gradio as gr
# Define the class names in the same order as during training
class_names = ["Defective tire", "Good condition tire", "Not a tire"]
def preprocess_image(image):
"""
Preprocess the uploaded image for prediction.
Args:
image (PIL.Image): Uploaded image from Gradio.
Returns:
np.ndarray: Preprocessed image tensor.
"""
try:
# Convert PIL image to array
img = image.resize((224, 224)) # Resize to match training dimensions
img_array = tf.keras.preprocessing.image.img_to_array(img) # Convert to NumPy array
# If your training used normalization, uncomment the next line
# img_array = img_array / 255.0
# Add batch dimension (1, 224, 224, 3)
img_array = np.expand_dims(img_array, axis=0)
return img_array
except Exception as e:
raise ValueError(f"Error preprocessing the image: {e}")
def predict_image(image):
"""
Predict the class of the uploaded image using the trained model.
Args:
image (PIL.Image): Uploaded image from Gradio.
Returns:
str: Predicted class and confidence score.
"""
try:
# Preprocess the image
img_array = preprocess_image(image)
# Make predictions
predictions = model.predict(img_array)
# Extract predicted class and confidence
predicted_class = class_names[np.argmax(predictions[0])]
confidence = np.max(predictions[0])
return f"Predicted Class: {predicted_class}, Confidence: {confidence:.2f}"
except Exception as e:
return f"Error during prediction: {e}"
# Load the trained model
model_path = "Class_ENB0.keras" # Replace with your model's file name
model = tf.keras.models.load_model(model_path)
# Define the Gradio interface
interface = gr.Interface(
fn=predict_image,
inputs=gr.Image(type="pil", label="Upload an Image"), # Accept image uploads
outputs=gr.Textbox(label="Prediction"), # Display prediction results
title="Tire Condition Classifier",
description="Upload an image of a tire to classify its condition."
)
# Launch the Gradio app
if __name__ == "__main__":
interface.launch()
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1" # Disable GPU usage