import streamlit as st
import tensorflow
from tensorflow.keras.models import load_model
import numpy as np
from PIL import Image
import pandas as pd
import os

# Load the saved TensorFlow model
model = load_model('traffic-sign-detection-model3.h5')


inputBasePath = 'D:\\traffic_Data\\'
path = 'D:\\traffic_Data\\DATA'
testingFolder = 'D:\\traffic_Data\\TEST'
classes = pd.read_csv('labels.csv')


# Function to preprocess the image
def preprocess_image(image):
    # Preprocess the image as required for your model
    # (e.g., resize, normalize pixel values)
    resized_image = image.resize((100,100))
    preprocessed_image = np.array(resized_image) / 255.0  # Normalize pixel values
    return preprocessed_image

# Function to make predictions
def predict(image):
    preprocessed_image = preprocess_image(image)
    prediction = model.predict(np.expand_dims(preprocessed_image, axis=0))
    return prediction

# Streamlit app
def main():
    st.title('Traffic Sign Detection')

    uploaded_image = st.file_uploader("Upload Image", type=['jpg', 'png', 'jpeg'])

    if uploaded_image is not None:
        # Display the uploaded image
        image = Image.open(uploaded_image)
        st.image(image, caption='Uploaded Image', use_column_width=True)

       # Predict button
        if st.button('Predict'):
            # Make prediction
            prediction = predict(image)
            predicted_class = np.argmax(prediction, axis=1)
            #st.write(predicted_class)
            class_mapping = dict(zip(classes['ClassId'], classes['Name']))
            predicted_label = class_mapping.get(predicted_class[0])
            # st.write(predicted_label)
            # st.write(predicted_class)
            
            # Display prediction result
            
            st.write('Prediction:', predicted_label)

if __name__ == '__main__':
    main()