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