import gradio as gr
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import load_img, img_to_array

# Load the trained model
model = load_model('model.h5')

examples = ['00000.png','00001.png','00086.png','00136.png','00011.png','00028.png','00039.png']


# Define class names (replace with your actual class names)
class_names = [       "Speed limit 20 km/h",
    "Speed limit 30 km/h",
    "Speed limit 50 km/h",
    "Speed limit 60 km/h",
    "Speed limit 70 km/h",
    "Speed limit 80 km/h",
    "End of speed limit 80 km/h",
    "Speed limit 100 km/h",
    "Speed limit 120 km/h",
    "No passing",
    "No passing for vehicles over 3.5 metric tons",
    "Right-of-way at the next intersection",
    "Priority road",
    "Yield",
    "Stop",
    "No vehicles",
    "Vehicles over 3.5 metric tons prohibited",
    "No entry",
    "General caution",
    "Dangerous curve to the left",
    "Dangerous curve to the right",
    "Double curve",
    "Bumpy road",
    "Slippery road",
    "Road narrows on the right",
    "Road work",
    "Traffic signals",
    "Pedestrians",
    "Children crossing",
    "Bicycles crossing",
    "Beware of ice/snow",
    "Wild animals crossing",
    "End of all speed and passing limits",
    "Turn right ahead",
    "Turn left ahead",
    "Ahead only",
    "Go straight or right",
    "Go straight or left",
    "Keep right",
    "Keep left",
    "Roundabout mandatory",
    "End of no passing",
    "End of no passing by vehicles over 3.5 metric tons"]  # Update this with your classes

# Preprocessing function
def preprocess_image(image):
    if image.mode != 'RGB':
        image = image.convert('RGB') 
    image = image.resize((30, 30))  # Resize to model's input size
    image_array = img_to_array(image) / 255.0  # Normalize to [0, 1]
    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
    return image_array

# Prediction function
def classify_image(image):
    # Preprocess the image
    preprocessed_image = preprocess_image(image)
    
    # Make prediction
    predictions = model.predict(preprocessed_image)
    
    # Get the class with the highest probability
    predicted_class = np.argmax(predictions, axis=1)[0]
    
    # Get the class name
    class_name = class_names[predicted_class]
    
    return f"Predicted Class: {class_name}"

# Create Gradio interface
interface = gr.Interface(
    fn=classify_image,  # Function to handle predictions
    inputs=gr.Image(type="pil", label="Original Images"),  # Accept image input
    outputs=gr.Textbox(label="Label "),  # Display output as text
    title="German Traffic Sign Recognition",
    description="Upload an image to classify it into one of the predefined classes.",
    examples=examples
)

# Launch the app
if __name__ == "__main__":
    interface.launch()