import streamlit as st
import tensorflow as tf
import numpy as np
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from PIL import Image

# Load the trained model
MODEL_PATH = "image_model.h5"
model = tf.keras.models.load_model(MODEL_PATH)

# Define image dimensions
IMG_WIDTH, IMG_HEIGHT = 64, 48

# Class labels with descriptions
CLASS_LABELS = {
    0: "Normal driving (safe)",
    1: "Texting - right",
    2: "Talking on the phone - right",
    3: "Texting - left",
    4: "Talking on the phone - left",
    5: "Operating the radio",
    6: "Drinking",
    7: "Reaching behind",
    8: "Hair and makeup",
    9: "Talking to passenger"
}

def predict_image(image):
    # Preprocess the image
    img_array = img_to_array(image)
    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
    img_array = img_array / 255.0  # Normalize pixel values to [0, 1]

    # Make a prediction
    predictions = model.predict(img_array)
    predicted_class = np.argmax(predictions, axis=1)[0]
    confidence = np.max(predictions)

    return CLASS_LABELS[predicted_class], confidence

# Streamlit app
st.title("Driver Distraction Detection")
st.markdown("Team18 Image Project : Sayandip Bhattacharyya, Purnendu Rudrapal, Sridatta Das, Sidhartha Karjee")

# File upload
uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Display the uploaded image
    image = Image.open(uploaded_file).convert("RGB")
    resized_image = image.resize((IMG_WIDTH, IMG_HEIGHT))
    st.image(image, caption="Uploaded Image", use_container_width=True)

    # Predict the class
    with st.spinner("Predicting..."):
        predicted_class, confidence = predict_image(resized_image)

    # Display the result
    st.subheader("Prediction")
    st.write(f"**Class:** {predicted_class}")
    st.write(f"**Confidence:** {confidence:.2%}")